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

做这道题之前，先完成1685. 有序数组中差绝对值之和 - 力扣（LeetCode）

 一般性的，我们能在这类题目中总结出以下规律：

求解有序数组中每个元素与q的差值res时，q 与数组交于 j ：

left = q*j - pre[j]
right = pre[n] - pre[j] - q * (n - j)
res = left + right

那么假设我们已经求得a，a中存储的是nums中相同元素的下标，比如nums=[1,3,1,1,2]，则对于元素1，它的a = [0, 2, 3]。

求解的结果存在res中，对元素1，我们需要求解res[0], res[2], res[3]

res[0] =  |0 - 0| - |0 - 2| + |0 - 3|

res[2] =  |0 - 2| - |2 - 2| + |2 - 3|

res[3] =  |0 - 3| - |2 - 3| + |3 - 3|

for j, q in enumerate(a):
    ...
    res[j] = left + right

其中：

left = q * j - pre[j]
right = pre[n] - pre[j] - q * (n - j)

可以看到就本题的核心就和1685完全一致，问题就只剩下了怎么构造这个数组a，这里我用了哈希表存储：

groups = defaultdict(list)
for i, x in enumerate(nums):
    groups[x].append(i)  # 记录相同元素的下标

完整代码如下

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))  # 计算数组a前i个元素的前缀和，以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

类似题目：(31条消息) 前缀和：leetcode2602数组元素全部相等的最少操作次数_坠金的博客-CSDN博客