2615. 等值距离和

题目

2615. 等值距离和 - 力扣（LeetCode）

思路

相同元素分组+前缀和
将相同元素的角标放到一组中，计算每组角标的前缀和。遍历每组角标，按题目要求求和。

代码

角标分组的复杂度可优化到 $O(n)$

func distance(nums []int) []int64 {
	ind := make([]int, len(nums))
	arr := make([]int64, len(nums))
	preSum := make([]int64, len(nums)+1)
	for i := range nums {
		ind[i] = i
	}
	sort.Slice(ind, func(i, j int) bool {
        if nums[ind[i]] == nums[ind[j]]{
            return ind[i] < ind[j]
        }
		return nums[ind[i]] < nums[ind[j]]
	})
	//fmt.Println(ind)
	preSum[0] = 0
	for i := range nums {
		preSum[i+1] = preSum[i] + int64(ind[i])
	}
	for l := 0; l < len(nums); {
		r := sort.Search(len(ind), func(i int) bool {
			return nums[ind[i]] > nums[ind[l]]
		})
		for i := l; i < r; i++ {
			arr[ind[i]] += int64((i-l)*ind[i]) - (preSum[i] - preSum[l])
			arr[ind[i]] += preSum[r] - preSum[i] - int64((r-i)*ind[i])
        }
		l = r
	}
	return arr
}