leetcode之前缀和

最新推荐文章于 2024-07-14 02:25:25 发布

HarryXxc

最新推荐文章于 2024-07-14 02:25:25 发布

阅读量1.1k

点赞数 1

分类专栏：数据结构与算法文章标签：算法 leetcode 数据结构 java 动态规划

本文链接：https://blog.csdn.net/harryshumxu/article/details/106385290

版权

数据结构与算法专栏收录该内容

24 篇文章 16 订阅

订阅专栏

在动态规划中我们经常会用到前缀和这个方法，我们先来简单看个例子。

借用leetcode第560题：和为k的子数组：给定一个整数数组和一个整数 k，你需要找到该数组中和为 k 的连续的子数组的个数。
示例：
输入:nums = [1,1,1], k = 2
输出: 2 , [1,1] 与 [1,1] 为两种不同的情况。

其实很容易想到的是暴力求解，但是时间复杂度为O( $n^2$ )，所以我们这里直接讲解前缀和；
我们遍历每一个数组，定义一个 preSum[i] 将数组下标在[0…i]之间的和保存起来，我们只要能够找到满足以下公式的 preSum[i] 就可以了。 $p r e S u m [i] - p r e S u m [j] = k$

public class Solution {
    public int subarraySum(int[] nums, int k) {
        int len = nums.length;
        // 计算前缀和数组
        int[] preSum = new int[len + 1];
        preSum[0] = 0;
        for (int i = 0; i < len; i++) {
            preSum[i + 1] = preSum[i] + nums[i];
        }
        int count = 0;
        for (int left = 0; left < len; left++) {
            for (int right = left; right < len; right++) {
                if (preSum[right + 1] - preSum[left] == k) {
                    count++;
                }
            }
        }
        return count;
    }
}

但是如上写出代码之后，发现这个怎么还比不上暴力求解写的舒服呢？？？别急，我们现在就开始优化这个前缀和。因为我们刚才说到找到的preSum[i] 需要满足： $p r e S u m [i] - p r e S u m [j] = k$
那我们改变一下公式的写法： $p r e S u m [j] = p r e S u m [i] - k$
这样看起来是不是更加清楚一些？因为我们并不需要关心究竟是哪一个数组下标满足以上的公式，只需要知道前缀和 preSum[j] 满足当前下标的前缀和 preSum[i] - k 的数量，所以我们这里在这里可以使用HashMap。

class Solution {
    public int subarraySum(int[] nums, int k) {
        //前缀和优化
        int count = 0;
        int pre = 0;
        HashMap<Integer,Integer> map = new HashMap();
        map.put(0,1);
        for(int i=0;i<nums.length;i++){
            pre += nums[i];
            if(map.containsKey(pre - k)){
                count += map.get(pre-k);
            }
            //这里map.getOrDefault(pre, 0)的作用是如果map函数中有pre这个key，那么就用这个key对应的value值，如果没有就使用默认值0
            map.put(pre, map.getOrDefault(pre, 0) + 1);
        }
        return count;
    }
}

看到这里我们再来看一道题：给定一个整数数组 A，返回其中元素之和可被 K 整除的（连续、非空）子数组的数目。
示例：
输入：A = [4,5,0,-2,-3,1], K = 5
输出：7
解释：
有 7 个子数组满足其元素之和可被 K = 5 整除：
[4, 5, 0, -2, -3, 1], [5], [5, 0], [5, 0, -2, -3], [0], [0, -2, -3], [-2, -3]

这道题同样可以使用前缀和+hashmap优化的方式去做，首先我们要确定找到满足以下公式的前缀和 $\ mod \ K= 0$
同样的，我们这里再做一个转换： $\ mod \ K = preSum[j]\ mod\ K$
因为我们也并不需要关心究竟是哪一个数组下标满足以上的公式，只需要知道前缀和 preSum[j] % K 满足当前下标的前缀和 preSum[i] % K 的数量。但是这里我们需要注意，因为需要用到求模（会出现负数），所以需要对求模做一些调整。举一个简单的例子

preSum[i] = 4, preSum[j] = -3, K = 7;
(preSum[i] -preSum[j]) % K = (4-(-3)) % 7 = 0;满足上式；
转换后得到的是：4 % 7 = 4 和 -3 % 7 = -3 ，这两个这样看起来明显不等，但是他们应该是在一组内。

所以我们这里需要做以下调整： $\ mod\ K + K) \mod\ K$

class Solution {
    public int subarraysDivByK(int[] A, int K) {
        //前缀和
        int count = 0;
        int len = A.length;
        int preSum = 0;
        HashMap<Integer,Integer> map = new HashMap();
        map.put(0,1);
        for(int i=0;i<len;i++){
            preSum += A[i];
            int temp = (preSum % K + K) % K;
            if(map.containsKey(temp)){
                count += map.get(temp);
            }
            map.put(temp, map.getOrDefault(temp, 0) + 1);
        }
        return count;
    }
}

这里又考虑到我们的前缀和求模之后的数组下标在[0…k-1]之间，所以我们可以用数组来替代HashMap。

class Solution {
    public int subarraysDivByK(int[] A, int K) {
        //前缀和
        int count = 0;
        int len = A.length;
        int preSum = 0;
        int[] map = new int[K];//因为mod之后的元素只能在[0,k-1]之间
        map[0] = 1;
        for(int i=0;i<len;i++){
            preSum += A[i];
            int temp = (preSum % K + K) % K;
            count += map[temp];
            map[temp] += 1;
        }
        return count;
    }
}