本文属于「征服LeetCode」系列文章之一，这一系列正式开始于2021/08/12。由于LeetCode上部分题目有锁，本系列将至少持续到刷完所有无锁题之日为止；由于LeetCode还在不断地创建新题，本系列的终止日期可能是永远。在这一系列刷题文章中，我不仅会讲解多种解题思路及其优化，还会用多种编程语言实现题解，涉及到通用解法时更将归纳总结出相应的算法模板。
为了方便在PC上运行调试、分享代码文件，我还建立了相关的仓库：https://github.com/memcpy0/LeetCode-Conquest。在这一仓库中，你不仅可以看到LeetCode原题链接、题解代码、题解文章链接、同类题目归纳、通用解法总结等，还可以看到原题出现频率和相关企业等重要信息。如果有其他优选题解，还可以一同分享给他人。
由于本系列文章的内容随时可能发生更新变动，欢迎关注和收藏征服LeetCode系列文章目录一文以作备忘。
给你一个下标从 0 开始的整数数组 nums 。现有一个长度等于 nums.length 的数组 arr 。对于满足 nums[j] == nums[i] 且 j != i 的所有 j ，arr[i] 等于所有 |i - j| 之和。如果不存在这样的 j ，则令 arr[i] 等于 0 。
返回数组 arr 。
示例 1：
java
输入：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 的其他下标。示例 2：
java
输入：nums = [0,5,3]
输出：[0,0,0]
解释：因为 nums 中的元素互不相同，对于所有 i ，都有 arr[i] = 0 。提示：
1 <= nums.length <= 10^50 <= nums[i] <= 10^9
方法1 相同元素分组+前缀和
按照相同元素将下标分组后，再计算。思路和[[2602. 使数组元素全部相等的最少操作次数]]是一样的，只是2602题是排序（相同元素自然就聚集在一起）。但由于目标位置就是数组中的下标，所以无需二分。
java
class Solution {
public long[] distance(int[] nums) {
int n = nums.length;
var groups = new HashMap<Integer, List<Integer>>();
for (int i = 0; i < n; ++i) // 相同元素分到同一组，记录下标
groups.computeIfAbsent(nums[i], k -> new ArrayList<>()).add(i);
var ans = new long[n];
var s = new long[n + 1];
for (var a : groups.values()) {
int m = a.size();
for (int i = 0; i < m; ++i)
s[i + 1] = s[i] + a.get(i); // 这组元素下标的前缀和
for (int i = 0; i < m; ++i) {
int target = a.get(i);
long left = (long) target * i - s[i]; // 蓝色面积
long right = s[m] - s[i] - (long) target * (m - i); // 绿色面积
ans[target] = left + right;
}
}
return ans;
}
}
cpp
class Solution {
public:
vector<long long> distance(vector<int>& nums) {
int n = nums.size();
unordered_map<int, vector<int>> groups;
for (int i = 0; i < n; ++i)
groups[nums[i]].push_back(i); // 相同元素分到同一组，记录下标
vector<long long> ans(n);
long long s[n + 1];
s[0] = 0;
for (auto& [_, a]: groups) {
int m = a.size();
for (int i = 0; i < m; ++i)
s[i + 1] = s[i] + a[i]; // 前缀和
for (int i = 0; i < m; ++i) {
long long target = a[i];
long long left = target * i - s[i]; // 蓝色面积
long long right = s[m] - s[i] - target * (m - i); // 绿色面积
ans[target] = left + right;
}
}
return ans;
}
};
python
class Solution:
def distance(self, nums: List[int]) -> List[int]:
groups = defaultdict(list)
for i, x in enumerate(nums):
groups[x].append(i) # 相同元素分到同一组，记录下标
ans = [0] * len(nums)
for a in groups.values():
n = len(a)
s = list(accumulate(a, initial=0)) # 前缀和
for j, target in enumerate(a):
left = target * j - s[j] # 蓝色面积
right = s[n] - s[j] - target * (n - j) # 绿色面积
ans[target] = left + right
return ans
go
func distance(nums []int) []int64 {
groups := map[int][]int{}
for i, x := range nums {
groups[x] = append(groups[x], i) // 相同元素分到同一组，记录下标
}
ans := make([]int64, len(nums))
for _, a := range groups {
n := len(a)
s := make([]int, n+1)
for i, x := range a {
s[i+1] = s[i] + x // 前缀和
}
for i, target := range a {
left := target*i - s[i] // 蓝色面积
right := s[n] - s[i] - target*(n-i) // 绿色面积
ans[target] = int64(left + right)
}
}
return ans
}复杂度分析：
- 时间复杂度： O ( n ) O(n) O(n) ，其中 n n n 为 n u m s nums nums 的长度。
- 空间复杂度： O ( n ) O(n) O(n) 。
方法2 相同元素分组+考虑增量
分组后，对于其中一个组 a a a ，我们先暴力计算出 a [ 0 ] a[0] a[0] 到其他元素的距离之和，设为 s 0 s_0 s0 。然后利用 s 0 s_0 s0 ， O ( 1 ) O(1) O(1) 时间计算出 a [ 1 ] a[1] a[1] 到其他元素的距离之和。注意，前提是数组 a a a 已经从小到大排序。
我们不再暴力计算，而是思考：从 s 0 s_0 s0 到 s 1 s_1 s1 ，这个增量 s 1 − s 0 s_1 - s_0 s1−s0 是多少（可以是负数）？
设 n n n 为 a a a 的长度，从 s 0 s_0 s0 到 s 1 s_1 s1 ，有 1 1 1 个数的距离变大了 a [ 1 ] − a [ 0 ] a[1] - a[0] a[1]−a[0]（原来是到 a [ 0 ] a[0] a[0] 的距离，现在是到 a [ 1 ] a[1] a[1] 的距离；原来是 a [ 0 ] a[0] a[0] 到 a [ 0 ] a[0] a[0] 的距离为 0 0 0 ，现在是 a [ 0 ] a[0] a[0] 到 a [ 1 ] a[1] a[1] 的距离为 a [ 1 ] − a [ 0 ] a[1] - a[0] a[1]−a[0] ，其余 n − 1 n - 1 n−1 个数的距离变小了 a [ 1 ] − a [ 0 ] a[1] - a[0] a[1]−a[0]（原来是到 a [ 0 ] a[0] a[0] 的距离，现在是到 a [ 1 ] a[1] a[1] 的距离）。
所以对于 s 1 s_1 s1 ，我们可以 O ( 1 ) O(1) O(1) 知道 a [ 1 ] a[1] a[1] 到其他元素的距离之和为： s 0 + ( 2 − n ) ⋅ ( a [ 1 ] − a [ 0 ] ) s_0 + (2 - n) \cdot (a[1] - a[0]) s0+(2−n)⋅(a[1]−a[0])一般地，设 a [ i − 1 ] a[i - 1] a[i−1] 到其他元素的距离之和为 s s s ，那么 a [ i ] a[i] a[i] 到其他元素的距离之和 s i s_i si 为：
s i − 1 + ( 2 i − n ) ⋅ ( a [ i ] − a [ i − 1 ] ) s_{i-1} + (2i - n) \cdot (a[i] - a[i - 1]) si−1+(2i−n)⋅(a[i]−a[i−1])
- 在 a [ i ] a[i] a[i] 左侧的元素（包括 a [ i − 1 ] a[i - 1] a[i−1] ），共有 i i i 个元素。这些元素到当前目标的距离都变大 了，增加的量正好是 a [ i ] − a [ i − 1 ] a[i] - a[i - 1] a[i]−a[i−1] 。
- 在 a [ i ] a[i] a[i] 右侧的元素（包括 a [ i ] a[i] a[i] ），共有 n − i n -i n−i 个元素。这些元素到当前目标的距离都变小 了，减少的量也是 a [ i ] − a [ i − 1 ] a[i] - a[i - 1] a[i]−a[i−1] 。
- 于是总的变化量是： s i = s i − 1 + i ⋅ ( a [ i ] − a [ i − 1 ] ) − ( n − i ) ⋅ ( a [ i ] − a [ i − 1 ] ) = s i − 1 + ( 2 i − n ) ( a [ i ] − a [ i − 1 ] ) s_i = s_{i-1} + i\cdot (a[i] - a[i - 1]) - (n - i) \cdot (a[i] - a[i - 1]) = s_{i-1} + (2i -n) (a[i] - a[i-1]) si=si−1+i⋅(a[i]−a[i−1])−(n−i)⋅(a[i]−a[i−1])=si−1+(2i−n)(a[i]−a[i−1])
java
class Solution {
public long[] distance(int[] nums) {
int n = nums.length;
var groups = new HashMap<Integer, List<Integer>>();
for (int i = 0; i < n; ++i) // 相同元素分到同一组，记录下标
groups.computeIfAbsent(nums[i], k -> new ArrayList<>()).add(i);
var ans = new long[n];
for (var a : groups.values()) {
int m = a.size();
long s = 0;
for (int x : a)
s += x - a.get(0); // a[0] 到其它下标的距离之和
ans[a.get(0)] = s;
for (int i = 1; i < m; ++i)
// 从计算 a[i-1] 到计算 a[i]，考虑 s 增加了多少
ans[a.get(i)] = s += (long) (i * 2 - m) * (a.get(i) - a.get(i - 1));
}
return ans;
}
}
cpp
class Solution {
public:
vector<long long> distance(vector<int>& nums) {
int n = nums.size();
unordered_map<int, vector<int>> groups;
for (int i = 0; i < n; ++i)
groups[nums[i]].push_back(i); // 相同元素分到同一组，记录下标
vector<long long> ans(n);
for (auto& [_, a]: groups) {
int m = a.size();
long long s = 0;
for (int x: a)
s += x - a[0]; // a[0] 到其它下标的距离之和
ans[a[0]] = s;
for (int i = 1; i < m; ++i)
// 从计算 a[i-1] 到计算 a[i]，考虑 s 增加了多少
ans[a[i]] = s += (long long) (i * 2 - m) * (a[i] - a[i - 1]);
}
return ans;
}
};
python
class Solution:
def distance(self, nums: List[int]) -> List[int]:
groups = defaultdict(list)
for i, x in enumerate(nums):
groups[x].append(i) # 相同元素分到同一组，记录下标
ans = [0] * len(nums)
for a in groups.values():
n = len(a)
s = sum(x - a[0] for x in a) # a[0] 到其它下标的距离之和
ans[a[0]] = s
for i in range(1, n):
# 从计算 a[i-1] 到计算 a[i]，考虑 s 增加了多少
s += (i * 2 - n) * (a[i] - a[i - 1])
ans[a[i]] = s
return ans
go
func distance(nums []int) []int64 {
groups := map[int][]int{}
for i, x := range nums {
groups[x] = append(groups[x], i) // 相同元素分到同一组，记录下标
}
ans := make([]int64, len(nums))
for _, a := range groups {
n := len(a)
s := int64(0)
for _, x := range a {
s += int64(x - a[0]) // a[0] 到其它下标的距离之和
}
ans[a[0]] = s
for i := 1; i < n; i++ {
// 从计算 a[i-1] 到计算 a[i]，考虑 s 增加了多少
s += int64(i*2-n) * int64(a[i]-a[i-1])
ans[a[i]] = s
}
}
return ans
}复杂度分析：
- 时间复杂度： O ( n ) O(n) O(n) ，其中 n n n 为 n u m s nums nums 的长度。
- 空间复杂度： O ( n ) O(n) O(n) 。