力扣周赛第340场 6360. 等值距离和

题目描述

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

题目链接

示例

输入：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 的其他下标。

方法一：分组+前缀和

1. 使用Map表分组，将值一样的下标放在同一数组里面。
2. 遍历每组数组，并计算前缀和。
   前缀和为每个元素自身与它前面元素的和。例如：第i个元素的前缀和为nums[0] + nums[1] + nums[2] + … + nums[i].
3. 根据前缀和计算出各距离和。
   假设有一个数组为[1,2,1,3,1,1,2,4,1,1]
   1对应分组后结果：arr = [0,2,4,5,8,9]
   1对应前缀和后结果：ans = [0,2,6,11,19,28]
   当 i = 3时
   先计算下标为3元素前的和为(11-0)+(11-2)+(11-4)=11*3 - (0+2+4). 因为有了前缀和所以(0+2+4)可以直接用ans[2]来代替. 所以最终式子应该为arr[i] * i - ans[i-1].
   再计算下标为3元素后的和为(19-11)+(28-11) = (19+28) - 11 * 2.
   同理用前缀和去替换(19+28)，即为arr[arr.length - 1] - arr[i]. 因此最终式子应为(arr[arr.length - 1] - arr[i]) - arr[i] * (arr.length - 1 - i).

function distance(nums: number[]): number[] {
	// 分组
    const map: Map<number, number[]> = new Map();
    nums.forEach((num,idx)=> {
        const tmp = map.get(num);
        if(tmp) {
            tmp.push(idx);
        } else {
            map.set(num, [idx]);
        }
    })
    const res: number[] = new Array(nums.length).fill(0);
    for(const arr of map.values()) {
        const ans: number[] = [];
        let sum = 0;
        // 计算前缀和
        arr.forEach(num=> {
            sum += num;
            ans.push(sum);
        })
        const n = arr.length;
        // 根据前缀和算各距离和
        for(let i = 0;i < n;i++) {
            let front = arr[i]*i - (i == 0 ? 0 : ans[i-1]);
            let back = (ans[n-1] - ans[i]) - arr[i] * (n - i - 1);
            res[arr[i]] = front + back;
        }
    }
    return res;
};

时间复杂度：O(n)
空间复杂度：O(n)

方法二：分组+增量

1. 使用Map表分组，将值一样的下标放在同一数组里面。
2. 遍历Map表，先算出每组数组第一个元素的距离和。
3. 遍历每组数组，根据前一个的距离和算出当前的距离和。即计算当前距离和比上一个距离和增加了多少。
   假设有一个数组为[1,2,1,3,1,1,2,4,1,1]
   1对应分组后结果：arr = [0,2,4,5,8,9]
   我们可以看出相较于第一个元素0，2跟自己本身以及它后面的元素的距离都减少了(arr[1]-arr[0]) 。所以减少了n-1个(arr[1]-arr[0])的距离，即增加了(1-n)*(arr[1]-arr[0])。
   而2跟它前面的元素距离都增加了(arr[1]-arr[0]) ，即增加了1*(arr[1]-arr[0])。
   所以我们可以得到当前元素的相对于上一个元素的距离和增量应为(2*i-n)*(arr[1]-arr[0])。

function distance(nums: number[]): number[] {
	// 分组
    const map: Map<number, number[]> = new Map();
    nums.forEach((num,idx)=> {
        const tmp = map.get(num);
        if(tmp) {
            tmp.push(idx);
        } else {
            map.set(num, [idx]);
        }
    })
    const res: number[] = new Array(nums.length).fill(0);
    for(const arr of map.values()) {
        // 先算出第一个的距离和
        let frontNum = 0;
        for(let i = 1;i < arr.length;i++) {
            frontNum += arr[i] - arr[0];
        }
        res[arr[0]] = frontNum;
        const n = arr.length;
        // 根据前一个的距离和算出当前的距离和
        for(let i = 1; i<arr.length;i++) {
            frontNum += (2*i - n)*(arr[i] - arr[i-1]);
            res[arr[i]] = frontNum;
        }
    }
    return res;
};

时间复杂度：O(n)
空间复杂度：O(n)

参考链接

两种方法：照搬 2602 / 考虑距离之和的增量 O(n)