题目描述
给你一个下标从 0 开始的整数数组 nums
。现有一个长度等于 nums.length
的数组 arr
。对于满足 nums[j] == nums[i]
且 j != i
的所有 j
,arr[i]
等于所有 |i - j|
之和。如果不存在这样的 j
,则令 arr[i]
等于 0
。
返回数组 arr
。
样例
输入:nums = [1,3,1,1,2]
输出:[5,0,3,4,0]
解释:
i = 0,nums[0] == nums[2] 且 nums[0] == nums[3]。因此,arr[0] = |0 - 2| + |0 - 3| = 5。
i = 1,arr[1] = 0 因为不存在值等于 3 的其他下标。
i = 2,nums[2] == nums[0] 且 nums[2] == nums[3]。因此,arr[2] = |2 - 0| + |2 - 3| = 3。
i = 3,nums[3] == nums[0] 且 nums[3] == nums[2]。因此,arr[3] = |3 - 0| + |3 - 2| = 4。
i = 4,arr[4] = 0 因为不存在值等于 2 的其他下标。
输入:nums = [0,5,3]
输出:[0,0,0]
解释:因为 nums 中的元素互不相同,对于所有 i,都有 arr[i] = 0。
限制
1 <= nums.length <= 10^5
0 <= nums[i] <= 10^9
算法
(思维题) $O(n)$
- 使用哈希表记录每个数字的所有出现位置。
- 对于单个数字,初始时,记录所有位置到第一个位置的距离总和 $sum$,然后从第二个位置开始遍历,遍历时,每次修改的值为 $(2i - m) * (p_i - p_{i-1})$,这里 $i$ 的下标从 $0$ 开始。
时间复杂度
- 每个位置被遍历常数次,故时间复杂度为 $O(n)$。
空间复杂度
- 需要 $O(n)$ 的额外空间存储哈希表和答案。
C++ 代码
#define LL long long
class Solution {
public:
vector<LL> distance(vector<int>& nums) {
const int n = nums.size();
unordered_map<int, vector<int>> p;
for (int i = 0; i < n; i++)
p[nums[i]].push_back(i);
vector<LL> ans(n);
for (const auto &[_, v] : p) {
int m = v.size();
LL sum = 0;
for (int i = 1; i < m; i++)
sum += v[i] - v[0];
ans[v[0]] = sum;
for (int i = 1; i < m; i++) {
sum += (LL)(2 * i - m) * (v[i] - v[i - 1]);
ans[v[i]] = sum;
}
}
return ans;
}
};