每日一题力扣 2615.等值距离和哈希表前缀和 C++ 题解

文章目录

题目描述
思路简述
- 前缀和推导过程
代码实现
复杂度分析
踩坑记录

题目描述

示例 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⁹

思路简述

这道题的数据范围达到了 10 5 10^5 105，如果用暴力解法（对每个元素都遍历整个数组找相同值并计算距离和），时间复杂度会是 O ( n 2 ) O(n^2) O(n2)，肯定会超时。因此我们需要找一个更高效的方法。

优化的核心思路是分组预处理 + 前缀和 + 数学推导：

分组：暴力算法超时得原因之一是我们每次都要从头计算arr[i]，所以我们想法是先将所有相同元素下标放入一个数组中，我们可以用一个unordere_map<int,vector<int>>进行存储用哈希表把所有相同元素的下标存到一起。因为我们是顺序遍历数组的，所以每个组里的下标天然就是按升序排列的，这为后续的数学计算提供了便利。
数学推导：对于一组有序的下标，我们不需要每次都重新计算每个元素到其他元素的距离和，而是可以利用前一个位置的计算结果，通过前缀和快速推导出当前位置的距离和。

前缀和推导过程

假设对于某个数值 x x x，它在数组中的所有下标按顺序存放在数组 a a a 中，即 a = [ a 0 , a 1 , . . . , a k − 1 ] a = [a_0, a_1, ..., a_{k-1}] a=[a0,a1,...,ak−1]（其中 k k k 是该数值出现的次数）。现在我们要计算对于下标 a i a_i ai，它到其他所有下标的距离和 S i S_i Si。

因为数组 a a a 是有序的，所以：

当 j < i j < i j<i 时， a i ≥ a j a_i \geq a_j ai≥aj，绝对值可以直接去掉： ∣ a i − a j ∣ = a i − a j |a_i - a_j| = a_i - a_j ∣ai−aj∣=ai−aj
当 j > i j > i j>i 时， a j ≥ a i a_j \geq a_i aj≥ai，绝对值可以直接去掉： ∣ a i − a j ∣ = a j − a i |a_i - a_j| = a_j - a_i ∣ai−aj∣=aj−ai

因此，距离和 S i S_i Si 可以拆分为两部分：
S i = ∑ j = 0 i − 1 ( a i − a j ) + ∑ j = i + 1 k − 1 ( a j − a i ) S_i = \sum_{j=0}^{i-1} (a_i - a_j) + \sum_{j=i+1}^{k-1} (a_j - a_i) Si=j=0∑i−1(ai−aj)+j=i+1∑k−1(aj−ai)

我们分别计算这两个部分：

左边部分（ j < i j < i j<i） ：
∑ j = 0 i − 1 ( a i − a j ) = i ⋅ a i − ∑ j = 0 i − 1 a j \sum_{j=0}^{i-1} (a_i - a_j) = i \cdot a_i - \sum_{j=0}^{i-1} a_j j=0∑i−1(ai−aj)=i⋅ai−j=0∑i−1aj

这里 ∑ j = 0 i − 1 a j \sum_{j=0}^{i-1} a_j ∑j=0i−1aj 是前 i i i 个元素的前缀和，我们记为 p r e _ s u m i pre\_sum_i pre_sumi。
右边部分（ j > i j > i j>i） ：
∑ j = i + 1 k − 1 ( a j − a i ) = ∑ j = i + 1 k − 1 a j − ( k − 1 − i ) ⋅ a i \sum_{j=i+1}^{k-1} (a_j - a_i) = \sum_{j=i+1}^{k-1} a_j - (k - 1 - i) \cdot a_i j=i+1∑k−1(aj−ai)=j=i+1∑k−1aj−(k−1−i)⋅ai

其中 ∑ j = i + 1 k − 1 a j \sum_{j=i+1}^{k-1} a_j ∑j=i+1k−1aj 等于总和减去前 i + 1 i+1 i+1 个元素的和，即 t o t a l _ s u m − p r e _ s u m i − a i total\_sum - pre\_sum_i - a_i total_sum−pre_sumi−ai。

将两部分合并：
S i = ( i ⋅ a i − p r e _ s u m i ) + ( ( t o t a l _ s u m − p r e _ s u m i − a i ) − ( k − 1 − i ) ⋅ a i ) = t o t a l _ s u m − 2 ⋅ p r e _ s u m i + i ⋅ a i − a i − ( k − 1 − i ) ⋅ a i = t o t a l _ s u m − 2 ⋅ p r e _ s u m i + a i ⋅ ( 2 i − k ) \begin{align*} S_i &= \left( i \cdot a_i - pre\_sum_i \right) + \left( (total\_sum - pre\_sum_i - a_i) - (k - 1 - i) \cdot a_i \right) \\ &= total\_sum - 2 \cdot pre\_sum_i + i \cdot a_i - a_i - (k - 1 - i) \cdot a_i \\ &= total\_sum - 2 \cdot pre\_sum_i + a_i \cdot (2i - k) \end{align*} Si=(i⋅ai−pre_sumi)+((total_sum−pre_sumi−ai)−(k−1−i)⋅ai)=total_sum−2⋅pre_sumi+i⋅ai−ai−(k−1−i)⋅ai=total_sum−2⋅pre_sumi+ai⋅(2i−k)

有了这个公式，我们就可以在遍历下标数组时，一边维护前缀和，一边 O ( 1 ) O(1) O(1) 地计算出每个位置的距离和。

代码实现

cpp 复制代码

class Solution {
public:
    vector<long long> distance(vector<int>& nums) {
        // 1. 分组：用哈希表存储相同元素的所有下标
        // key: 数组中的数值，value: 该数值对应的所有下标（按遍历顺序存储，天然有序）
        unordered_map<long long, vector<int>> hash;
        for(int i = 0; i < nums.size(); i++)
        {
            hash[nums[i]].push_back(i);
        }

        vector<long long> ret(nums.size(), 0); // 用于存储最终结果

        // 2. 遍历每一组相同的元素，计算距离和
        for(auto& [x, num] : hash)
        {
            int k = num.size(); // 当前组的元素个数
            if(k == 1) continue; // 如果只有一个元素，距离和为0，跳过

            long long total_sum = accumulate(num.begin(), num.end(), 0LL); // 当前组所有下标的总和
            long long pre_sum = 0; // 前缀和，维护前 i 个下标的和

            for(int i = 0; i < k; i++)
            {
                // 使用推导的公式计算当前下标 num[i] 的距离和
                // total_sum: 所有下标的总和
                // pre_sum: 前 i 个下标的和 (即 num[0] 到 num[i-1] 的和)
                // k: 组内元素个数
                ret[num[i]] = total_sum - 2 * pre_sum + (long long)num[i] * (2 * i - k);
                
                // 更新前缀和，为下一次循环做准备
                pre_sum += num[i];
            }
        }

        return ret;
    }
};

复杂度分析

时间复杂度 ： O ( n ) O(n) O(n)
- 第一步遍历数组存入哈希表： O ( n ) O(n) O(n)
- 第二步遍历哈希表：虽然是两层循环，但每个下标只会被处理一次，因此总时间也是 O ( n ) O(n) O(n)
空间复杂度 ： O ( n ) O(n) O(n)
- 哈希表需要存储所有 n n n 个下标： O ( n ) O(n) O(n)
- 结果数组需要 O ( n ) O(n) O(n) 空间

踩坑记录

数学推导是关键：一句话败在了数学，看着有序的数组感觉到了应该是数学题（但是直接看的题解得公式了）。
数据类型溢出 ：距离和可能会很大，还有 num[i] * (2 * i - k) 这部分乘法也容易溢出 int，所以一定要用 long long 来存储中间结果和最终结果。

如果这篇博客对你有帮助，别忘了点赞支持一下～也可以收藏起来，方便后续刷题复习时随时翻看。要是能顺手点个关注，爱弥斯还能得到漂泊者批准的游戏时间哦！

每日一题 力扣 2615.等值距离和 哈希表 前缀和 C++ 题解