leetcode前缀和(2615. 等值距离和)

原创已于 2024-08-26 23:41:03 修改 · 468 阅读

4 ·

CC 4.0 BY-SA版权

文章标签：

#算法 #java #数据结构

于 2024-08-26 23:38:01 首次发布

前言

经过前期的基础训练以及部分实战练习，粗略掌握了各种题型的解题思路。现阶段开始专项练习。

描述

给你一个下标从 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 <= 105
0 <= nums[i] <= 109

实现原理与步骤

前缀和遍历：从左到右遍历数组，对每个元素，计算它与前面所有相同元素的距离和。通过累积前缀和和相同元素的个数，可以高效计算这些距离。
后缀和遍历：从右到左再遍历一次数组，对每个元素，计算它与后面所有相同元素的距离和。类似地，通过累积后缀和和相同元素的个数进行计算。

实现代码

class Solution {
    public long[] distance(int[] nums) {
        int n = nums.length;
        long[] result = new long[n];

        // 前缀和部分
        //存储每个元素在遍历过程中出现的次数
        HashMap<Integer, Long> countMap = new HashMap<>();
        //存储每个元素在遍历过程中到达当前索引的距离和
        HashMap<Integer, Long> prefixSumMap = new HashMap<>();

        for (int i = 0; i < n; i++) {
            int num = nums[i];
            if (countMap.containsKey(num)) {
                //countMap.get(num) * i是计算从0开始的距离和
                //prefixSumMap.get(num)是之前所有相同元素的实际索引之和
                result[i] += countMap.get(num) * i - prefixSumMap.get(num);
            }

            countMap.put(num, countMap.getOrDefault(num, 0l) + 1);
            prefixSumMap.put(num, prefixSumMap.getOrDefault(num, 0l) + i);
        }

        // 重置哈希表
        countMap.clear();
        prefixSumMap.clear();

        // 后缀和部分
        for (int i = n - 1; i >= 0; i--) {
            int num = nums[i];
            if (countMap.containsKey(num)) {
                result[i] += prefixSumMap.get(num) - countMap.get(num) * i;
            }

            countMap.put(num, countMap.getOrDefault(num, 0l) + 1);
            prefixSumMap.put(num, prefixSumMap.getOrDefault(num, 0l) + i);
        }

        return result;
    }
}