等值距离和

难度: Medium

给你一个下标从 0 开始的整数数组 nums 。现有一个长度等于 nums.length 的数组 arr 。对于满足 nums[j] == nums[i] 且 j != i 的所有 j ，arr[i] 等于所有 |i - j| 之和。如果不存在这样的 j ，则令 arr[i] 等于 0 。

返回数组 arr 。

示例 1：

输入：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：

输入：nums = [0,5,3]
输出：[0,0,0]
解释：因为 nums 中的元素互不相同，对于所有 i ，都有 arr[i] = 0 。

提示：

1 <= nums.length <= 10⁵
0 <= nums[i] <= 10⁹

Submission

运行时间: 280 ms

内存: 50.5 MB

class Solution:
    def distance(self, nums: List[int]) -> List[int]:
        n = len(nums)
        ans, rec = [0] * n, collections.defaultdict(list)
        for i, num in enumerate(nums):
            rec[num].append(i)
        for ls in rec.values():
            if len(ls) > 1:
                pre, m = list(itertools.accumulate(ls, initial=0)), len(ls)
                for i, num in enumerate(ls):
                    ans[num] = num * i - pre[i] + pre[-1] - num * (m - i) - pre[i]

        return ans

Explain

该题解采用的是使用哈希表来记录每个数字及其出现的所有索引位置。接下来，对于哈希表中的每个条目，如果一个数字出现不止一次，则计算与该数字相等的所有数组元素之间的距离和。为了高效地计算距离和，使用累积和技巧。对于每个索引位置i，使用公式：\(ans[num] = num \times i - pre[i] + pre[-1] - num \times (m - i) - pre[i]\)，其中\(pre\)是索引位置的累积和数组，\(m\)是当前数字的出现次数。这种方法避免了对每个索引对的显式距离计算，从而提高了效率。

时间复杂度: O(n)

空间复杂度: O(n)

class Solution:
    def distance(self, nums: List[int]) -> List[int]:
        n = len(nums)
        ans, rec = [0] * n, collections.defaultdict(list)
        # 创建哈希表，键为数字值，值为索引列表
        for i, num in enumerate(nums):
            rec[num].append(i)
        # 遍历哈希表中的所有索引列表
        for ls in rec.values():
            if len(ls) > 1:
                # 计算前缀和
                pre, m = list(itertools.accumulate(ls, initial=0)), len(ls)
                # 对于每个索引，计算距离和
                for i, num in enumerate(ls):
                    ans[num] = num * i - pre[i] + pre[-1] - num * (m - i) - pre[i]

        return ans

Explore

在构建哈希表时，确实考虑了数组中可能存在的相同数字。这通过使用哈希表的数据结构实现，其中键为数字本身，而值为一个列表，这个列表记录了该数字在数组中出现的所有索引位置。当在遍历数组时遇到相同的数字，我们将其索引追加到哈希表中该数字对应的列表中。这种方法能够确保即使数字重复出现，也能准确记录其所有位置，为后续的距离计算提供必要的信息。

在解题中，前缀和数组`pre`用于存储数字索引的累积和。具体来说，对于一个给定的数字及其在数组中的所有索引列表，`pre[i]`表示从列表开始到第i个索引之前的所有索引的和。这通过使用`itertools.accumulate`函数计算得到，它会自动累加给定列表中的元素，并生成一个新的列表，其中每个元素是到当前位置的元素累积和。

公式`num * i - pre[i] + pre[-1] - num * (m - i) - pre[i]`用于计算等值距离和。其中，`num * i`是当前索引乘以数字值，`pre[i]`是到当前索引之前所有索引的和，`pre[-1]`是所有索引的总和，`num * (m - i)`是从当前索引后所有位置的数字值乘以个数，`pre[i]`再次被减去，用于调整前缀和的计算。各项结合起来，可以有效地计算出从当前索引到其他所有相同数字的索引的距离和。

使用累积和技巧来计算距离和主要是为了提高计算效率。如果直接计算每对相同数字之间的距离和，需要双重循环，时间复杂度为O(n^2)，这在数字量大时非常低效。而通过使用前缀和，可以将每次的距离计算简化为几个数学运算，从而将时间复杂度降低到O(n)。因此，累积和不仅简化了计算过程，也显著提高了算法的执行效率。