【LeetCode Hot 100】#560 和为 K 的子数组：初见前缀和

题目描述

给你一个整数数组 nums 和一个整数 k ，请你统计并返回该数组中和为 k 的子数组的个数 。

子数组是数组中元素的连续非空序列。

示例 1：

输入：nums = [1,1,1], k = 2
输出：2

示例 2：

输入：nums = [1,2,3], k = 3
输出：2

提示：

• 1 <= nums.length <= 2 * 10^4
• -1000 <= nums[i] <= 1000
• -10^7 <= k <= 10^7

解法

暴力枚举

暴力解法是最简单的一种和最容易想到的一种解法，我们只需要通过两个 for循环 求出所有可能的子数组，然再判断其是否和等于 k 即可。其时间复杂度为 O(n^2) ，时间复杂度过大，很容易超时的。

前缀和+哈希

我们可以利用前缀和和哈希表将时间复杂度优化到 O(n) ，这样就不会超时了。

前缀和： 前缀和就是从数组开头到当前位置所有元素的累加和。

思路：

定义前缀和 pre[i] 表示数组从下标 0 到 i 的元素之和，即：

pre[i] = nums[0] + nums[1] + ... + nums[i]

那么，子数组 nums[j..i] 的和可以表示为：

pre[i] - pre[j-1]

我们想要求有多少个子数组的和等于 k，即：

pre[i] - pre[j-1] = k

移项得：

pre[j-1] = pre[i] - k

也就是说，对于当前的前缀和 pre[i]，我们只需要知道在它之前有多少个前缀和等于 pre[i] - k，就能确定有多少个子数组以 i 结尾且和为 k。

因此，我们可以用一个哈希表 mp 来记录每个前缀和出现的次数，键为前缀和，值为该前缀和出现的次数。

注意：

• 初始化 mp.put(0, 1)，表示前缀和为 0 出现了 1 次（即空前缀），这是为了处理从数组开头开始的子数组。
• 遍历数组，维护当前前缀和 pre。
• 计算 temp = pre - k，在哈希表中查找 temp 出现的次数，累加到结果 count 中。
• 然后将当前前缀和 pre 的出现次数加一，存入哈希表。
• 注意顺序：必须先查询再更新哈希表，否则当 k = 0 时，会把当前前缀和自己也算进去，导致多算。

具体过程：

以 nums = [1, 2, 3], k = 3 为例：

最终 count = 2，对应子数组 [1,2] 和 [3]，与答案一致。

解答

class Solution {
    public int subarraySum(int[] nums, int k) {
        Map<Integer, Integer> mp = new HashMap<>();
        int count = 0;
        int pre = 0;

        // 空前缀和为 0，出现 1 次
        mp.put(0, 1);

        for (int i = 0; i < nums.length; i++) {
            pre += nums[i];                     // 计算当前前缀和
            int temp = pre - k;                 // 需要查找的前缀和
            count += mp.getOrDefault(temp, 0);  // 累加出现次数
            mp.put(pre, mp.getOrDefault(pre, 0) + 1); // 更新当前前缀和的出现次数
        }

        return count;
    }
}

复杂度分析

• 时间复杂度：O(n) ，其中 n 为数组长度。只需遍历一次数组，哈希表的插入和查询操作平均时间复杂度为 O(1)。
• 空间复杂度：O(n) ，最坏情况下所有前缀和都不相同，哈希表需要存储 n 个键值对。

其他

mp.put(0, 1) 为什么是 1 而不是 0？

• 0 代表空前缀，即“从数组开头之前累加”的和。
• 它出现 1 次，是一个合法存在的前缀。
• 当 pre - k == 0 时，说明存在一个从下标 0 开始的子数组和为 k。如果没有这个 (0, 1)，这类子数组就会被漏掉。

举例：nums = [3], k = 3。
遍历时 pre = 3，pre - k = 0，如果 mp 中没有 0，count 就是 0，答案错误。