Leetcode每日一题 —— 2615. 等值距离和

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

2615. 等值距离和 - 给你一个下标从 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]...

思路

也就是对于每个下标 i 对应的数字 nums[i]，要找到与其相等的所有其他数字 nums[j]，累加 abs(i-j)。

题目输入规模可能很大，时间复杂度要控制在 O(n\log n) 以下。

咱们在意的是某个数字 nums[i] 在数组内出现的所有下标，但如果直接按数字分组维护列表，扫描所耗费的时间从总体上看来是难以接受的。

既要涉及求和，而且还需要尽量是线性时间复杂度，能猜到可能能用到前缀和思想。

对于一个下标 i，我可以计算前缀中所有 i-j 之和，再计算其后缀中所有 j-i 之和，这个位置的前缀加后缀就是最终结果 arr[i] 了。

怎么生成这样的前缀和呢？从计算某个位置的前缀 prefix[i] 来思考：

比如 [1,3,1,1,1,2] 这个数组，若要计算 prefix[4]

1. 首先 nums[4]=1，要从上一个前缀和处转移过来，我们可以用哈希表存上一个 1 出现的地方：3
2. prefix[4] 可转移自 prefix[3]，但是怎么从 prefix[3] 计算出 prefix[4]？
3. prefix[3] 等于 (3-0) + (3-2)。而 prefix[4] 则等于 (4-0) + (4-2) + (4-3)，如果把 prefix[3] 的部分代入进去可以发现是 prefix[3] + (4-3) + (4-3) + (4-3)，也就是在 prefix[3] 的基础上重复累加离上一个 1 出现的位置的距离多次。很显然这个次数就是 1 之前出现的次数。

因此在递推的时候可以用哈希表维护 nums[i] 上一次出现的下标以及 nums[i] 之前出现过的次数。递推得到前缀和和后缀和即可。

代码

写出来感觉还是挺直白的，应该还有很大可优化空间。

class Solution {
public:
    vector<long long> distance(vector<int>& nums) {
        int n=nums.size();
        vector<long long> res(n,0);
        // 也就是要找到所有相等元素的下标距离的和
        // 输入规模比较大，可以分组做前缀和后缀和
        vector<long long> prefix(n,0);
        // 某个数字 -> (其最后一次出现的下标, 这个数字出现的次数)
        unordered_map<int,pair<int,int>> iMap;
        for(int i=0;i<n;i++){
            int appearCnt=0;
            if(iMap.count(nums[i])>0){
                // 之前这个数字出现过
                // 比如 1 ... 1 1
                // 第三次出现的 1 的和前缀相当于 [前两个 1 的距离] + [第三次 1 和第二次 1 的距离] * [前面 1 出现的次数 = 2]
                int prevIdx=iMap[nums[i]].first;
                appearCnt=iMap[nums[i]].second;
                prefix[i]=prefix[prevIdx]+(i-prevIdx)*appearCnt;
                // cout<<"prefix["<<i<<"]="<<prefix[i]<<endl;
            }
            iMap[nums[i]]=make_pair(i,appearCnt+1);
        }
        // 接着边生成后缀边得到结果
        vector<long long> postfix(n,0);
        iMap.clear();
        for(int i=n-1;i>=0;i--){
            int appearCnt=0;
            if(iMap.count(nums[i])>0){
                int prevIdx=iMap[nums[i]].first;
                appearCnt=iMap[nums[i]].second;
                postfix[i]=postfix[prevIdx]+(prevIdx-i)*appearCnt;
            }
            res[i]=prefix[i]+postfix[i];
            iMap[nums[i]]=make_pair(i,appearCnt+1);
        }
        return res;
    }
};

1 个帖子 - 1 位参与者

阅读完整话题