LeetCode 560. 和为k的子数组

Offer 驾到，掘友接招！我正在参与2022春招打卡活动，点击查看活动详情。

题目：给定一个整数数组nums和整数k，要求统计并返回该数组中和为k的连续子数组的个数。

解题思路

因为题目要求的是返回连续子数组的和为k，那么最简单的思路就是暴力枚举所有的可能，这样每个可能都比对一下，最终必然可以得到正确的结果，此处暴力枚举采用的是双层for循环，如下：

public int subarraySum(int[] nums, int k) {
    int res = 0, len = nums.length;
    for(int i=0;i<len;i++){
        int sum = 0;
        for(int j=i;j<len;j++){
            sum += nums[j];
            if(sum == k) res++;
        }
    }
    return res;
}

时间复杂度为$O(n^2)$，空间复杂度为$O(1)$。最终实际耗时1500ms+。

思路进阶（前缀和）

前缀和数组preSum[i]是指i元素前面元素之和，如果元素i的索引为0，那么preSum[0]=0，但此时是存在一个元素的前缀是0的。

假设当前数组为[3, 1, -1, 0, 1]，那么由该数组元素形成的前缀和数组为[0, 3, 4, 3, 3, 4]如果此时k为1，我们可以发现前缀和数组中索引2位置减去索引1位置的元素刚好等于k，也即原数组中索引为1的前面元素之和减去索引为0的前面元素之和等于k，显然此时1是满足的。同理前缀和数组索引5的位置和前面的1, 3, 4也满足该等式。

更一般的理解就是 在i....j这一段序列中，如果存在j的前面元素和（包括自身）减去i的前面元素和（包括自身）等于k，则i+1 ... j这一段序列必定是满足的。

由此可得代码：

public int subarraySum(int[] nums, int k) {
    int res = 0, len = nums.length;
    // 前缀和数组 preSum[i] 表示第i个元素的前缀和
    int[] preSum = new int[len + 1];
    for(int i=0;i<len;i++){
        preSum[i+1] = preSum[i] + nums[i];
    }

    for(int i=0;i<len;i++){
        for(int j=i;j<len;j++){
            if(preSum[j+1] - preSum[i] == k) res++;
        }
    }
    return res;
}

上述代码虽然思路进阶了但时间复杂度一点没有减少，并且空间复杂度还加大了，最终耗时1500ms+。

前缀和 + 哈希表

我们可以对上述的方法进行优化，每次计算前缀和都将前缀和存入map中，这样当计算到新的前缀和，将此前缀和减去k，查看map中是否存在这样的前缀和，如果存在则此时的序列满足条件。

为什么这样算呢，可能还是有点难理解，假设我们此时存在一个序列[0, ... i,...j,...n]，k为1。如果此时计算到了j的前面所有元素之和（包括自身）这里假设为5，那么将preSum[j] - k就等于4，而此时在map中找到了i的前面所有元素之和刚好为4，则此时i+1 ... j这一段序列显然满足。并且可能在i的前面也存在元素之和为4的情况，因此使用map存储就避免了循环检查，可得代码如下：

public int subarraySum(int[] nums, int k) {
    int res = 0, preSum = 0;
    HashMap<Integer, Integer> map = new HashMap<>();
    map.put(0, 1);
    for(int i=0;i<nums.length;i++){
        preSum += nums[i];
        if(map.containsKey(preSum-k)){
            res += map.get(preSum-k);
        }
        map.put(preSum, map.getOrDefault(preSum, 0)+1);
    }
    return res;
}

时间复杂度和空间复杂度都为$O(n)$。