题目 #

You are given a 0-indexed integer array nums. There exists an array arr of length nums.length, where arr[i] is the sum of |i - j| over all j such that nums[j] == nums[i] and j != i. If there is no such j, set arr[i] to be 0.

Return the array arr.

Example 1:

Input: nums = [1,3,1,1,2]
Output: [5,0,3,4,0]
Explanation:
When i = 0, nums[0] == nums[2] and nums[0] == nums[3]. Therefore, arr[0] = |0 - 2| + |0 - 3| = 5.
When i = 1, arr[1] = 0 because there is no other index with value 3.
When i = 2, nums[2] == nums[0] and nums[2] == nums[3]. Therefore, arr[2] = |2 - 0| + |2 - 3| = 3.
When i = 3, nums[3] == nums[0] and nums[3] == nums[2]. Therefore, arr[3] = |3 - 0| + |3 - 2| = 4.
When i = 4, arr[4] = 0 because there is no other index with value 2.

Example 2:

Input: nums = [0,5,3]
Output: [0,0,0]
Explanation: Since each element in nums is distinct, arr[i] = 0 for all i.

Constraints:

$1 <= nums.length <= 10^5$
$0 <= nums[i] <= 10^9$

思路1 正序+倒序计算 #

分析 #

把索引的差值分成两部分，当前索引减去前面的和当前索引被后面的减去的两部分
那么只算作为被减数的话，只有当前减去前面的，每一个数都是前一个数的和加上 $已知相同的数的数量 \times 当前索引减去前一个索引的差值$
只算作为减数的部分，倒序更好计算，每一个结果都是前一个数的结果加上 $已知相同数的数量 \times 前一个计算的索引减去当前索引的差值$

代码 #

 1func distance(nums []int) []int64 {
 2	res := make([]int64, len(nums))
 3	n := len(nums)
 4	// key均为值
 5	inMap := make(map[int]int)    // 正序的索引
 6	cntMap := make(map[int]int)   // 正序的数量
 7	// 正序遍历
 8	for i, v := range nums {
 9		if j, ok := inMap[v]; ok {
10			// 前面有
11			res[i] = res[j] + int64(cntMap[v] * (i - j))
12		}
13		cntMap[v]++
14		inMap[v] = i
15	}
16	inMap = make(map[int]int)    // 倒序的索引
17	cntMap = make(map[int]int)   // 倒序的数量
18	sumMap := make(map[int]int64) // 倒序的和
19	for i := range nums {
20		i = n - 1 - i
21		v := nums[i]
22		if j, ok := inMap[v]; ok {
23			// 前面有
24			sumMap[v] += int64(cntMap[v] * (j - i))
25			res[i] += sumMap[v]
26		}
27		cntMap[v]++
28		inMap[v] = i
29	}
30	return res
31}

思路2 分组+前缀和 #

分析 #

要求的其实是上图的绿色面积的和
分成两段，都可以用前缀和与当前数乘以n来处理
那么第一步先分组，将相同数分到一起，然后计算前缀和，最后合并成结果

代码 #

 1func distance1(nums []int) []int64 {
 2	grp := make(map[int][]int) // 数组存放索引
 3	for i, v := range nums {
 4		t := grp[v]
 5		if t == nil {
 6			t = make([]int, 0, 1)
 7		}
 8		t = append(t, i)
 9		grp[v] = t
10	}
11
12	res := make([]int64, len(nums))
13	// 遍历分组计算结果
14	for _, v := range grp {
15		n := len(v)
16		if n == 1 {
17			continue
18		}
19		s := make([]int, n) // 计算前缀和
20		s[0] = v[0]
21		for i := 1; i < n; i++ {
22			s[i] = s[i-1] + v[i]
23		}
24		for i := range v {
25			res[v[i]] = int64((i+1)*v[i] - s[i] + (s[n-1] - s[i]) - (n-1-i)*v[i])
26		}
27	}
28	return res
29}

思路3 分组+数学分析 #

分析 #

分组不变，想一下对每一个分组v来说，如果有了所有数到第一个数的距离和a[0]，那么有

$$ a[1] = a[0] + (v[1] - v[0]) - (n-1) * (v[1] - v[0]) $$

推出

$$ a[i] = a[i-1] - (n-i-i) * (v[i]-v[i-1]) $$

代码 #

 1func distance2(nums []int) []int64 {
 2	grp := make(map[int][]int) // 数组存放索引
 3	for i, v := range nums {
 4		grp[v] = append(grp[v], i)
 5	}
 6
 7	res := make([]int64, len(nums))
 8	// 遍历分组计算结果
 9	for _, v := range grp {
10		n := len(v)
11		if n == 1 {
12			continue
13		}
14		var s int64 = 0
15		for _, v1 := range v[1:] {
16			s += int64(v1 - v[0])
17		}
18		res[v[0]] = s
19		for i := 1; i < n; i++ {
20			res[v[i]] = res[v[i-1]] - int64((n-i-i)*(v[i]-v[i-1]))
21		}
22	}
23	return res
24}