LeetCode 2615.等值距离和：分组(哈希表+前缀和)

【LetMeFly】2615.等值距离和：分组(哈希表+前缀和)

力扣题目链接：https://leetcode.cn/problems/sum-of-distances/

给你一个下标从 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⁹

解题方法：哈希表 + 前缀和

读起来题意理解稍难，在此转述下：

对于数组中下标为 i i i的元素 n u m s [ i ] nums[i] nums[i]，它对应的要返回的结果 a n s [ i ] ans[i] ans[i]怎么算呢？

对于 n u m s nums nums中所有和 n u m s [ i ] nums[i] nums[i]相等的 n u m s [ j ] nums[j] nums[j]们，计算 a b s ( j − i ) abs(j-i) abs(j−i)之和即为所求。

显而易见，对于元素 n u m s [ i ] nums[i] nums[i]，其最终结果仅与和 n u m s [ i ] nums[i] nums[i]相等的元素有关。因此我们可以使用一个哈希表，键为 n u m s [ i ] nums[i] nums[i]值为所有等于 n u m s [ i ] nums[i] nums[i]的下标们，遍历一遍即可实现原始数组按照值的分组。

好了，分好组了，现在我们得到了很多的 i d x s idxs idxs数组，每个 i d x s idxs idxs数组存放的都是在 n u m s nums nums中值相等的下标们。问题就变成了给你一个升序的 i d x s idxs idxs数组，对于其中的元素 a a a，求数组中其他所有元素和 a a a的的差的绝对值之和。

怎么算？前缀和+后缀和就好：

创建一个后缀和数组 s u f f i x suffix suffix，令 s u f f i x [ i ] suffix[i] suffix[i]为 ∑ j = i j < l e n ( i d x s ) i d x s [ j ] \sum_{j=i}^{j\lt len(idxs)} idxs[j] ∑j=ij<len(idxs)idxs[j]（即 i d x s idxs idxs数组中从 i d x s [ i ] idxs[i] idxs[i]到最后一个元素的元素和）。

从前到后遍历 i d x s idxs idxs数组，使用一个变量 p r e f i x prefix prefix， p r e f i x prefix prefix代表这个元素之前所有元素的和（不包含当前遍历到的元素）。

那么， i d x s [ i ] idxs[i] idxs[i]要算出的最终结果就是： i i i后面所有元素之和( s u f f i x [ i + 1 ] suffix[i+1] suffix[i+1]) − - − 后面元素个数 × i d x s [ i ] + i \times idxs[i] + i ×idxs[i]+i前面元素个数 × i d x s [ i ] \times idxs[i] ×idxs[i] − - − 前面元素之和 p r e f i x prefix prefix。

另外，也可以不使用 s u f f i x suffix suffix数组，改为 s u f f i x [ i + 1 ] suffix[i+1] suffix[i+1]由 t o t a l − p r e f i x − i d x s [ i ] total-prefix-idxs[i] total−prefix−idxs[i]得出。

时间复杂度 O ( l e n ( n u m s ) ) O(len(nums)) O(len(nums))
空间复杂度 O ( l e n ( n u m s ) ) O(len(nums)) O(len(nums))

AC代码

C++

cpp 复制代码

/*
 * @LastEditTime: 2026-04-24 13:53:47
 */
typedef long long ll;
class Solution {
private:
    void cal(vector<int>& idxs, vector<ll>& ans) {
        int n = idxs.size();
        ll total = 0;
        for (int i = n - 1; i >= 0; i--) {
            total += idxs[i];
        }
        ll prefix = 0;
        for (int i = 0; i < n; i++) {
            ans[idxs[i]] += (total - prefix - idxs[i]) - (ll)(n - i - 1) * idxs[i];
            ans[idxs[i]] += (ll)i * idxs[i] - prefix;
            prefix += idxs[i];
        }
    }
public:
    vector<ll> distance(vector<int>& nums) {
        unordered_map<int, vector<int>> ma;
        for (int i = 0; i < nums.size(); i++) {
            ma[nums[i]].push_back(i);
        }

        vector<ll> ans(nums.size());
        for (auto&& [_, idxs] : ma) {
            cal(idxs, ans);
        }
        return ans;
    }
};

同步发文于CSDN和我的个人博客，原创不易，转载经作者同意后请附上原文链接哦~

千篇源码题解已开源