和为K的子数组（LeetCode 560）

1.问题描述

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

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

示例 1：

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

示例 2：

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

示例 3：

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

提示：

• 1 <= nums.length <= 2 * 104
• -1000 <= nums[i] <= 1000
• -107 <= k <= 107

2.难度等级

Medium。

3.热门指数

★★★★☆

出题公司：阿里、腾讯、字节。

4.解题思路

方法一：枚举

最容易想到是暴力枚举。

考虑以 i 结尾和为 k 的连续子数组个数，我们需要统计符合条件的下标 j 的个数，其中 0≤j≤i 且 [j…i] 这个子数组的和恰好为 k 。

可能有读者会认为假定我们确定了子数组的开头和结尾，还需要 O(n) 的时间复杂度遍历子数组来求和，那样复杂度就将达到 O(n^3) 从而无法通过所有测试用例。但是如果我们知道 [j,i] 子数组的和，就能 O(1) 推出 [j−1,i] 的和，因此这部分的遍历求和是不需要的，我们在枚举下标 j 的时候已经能 O(1) 求出 [j,i] 的子数组之和。

时间复杂度： O(n^2)，其中 n 为数组的长度。枚举子数组开头和结尾需要 O(n^2) 的时间，其中求和需要 O(1) 的时间复杂度，因此总时间复杂度为 O(n^2)。

空间复杂度： O(1)。只需要常数空间存放若干变量。

下面以 Golang 为例给出实现。

func subarraySum(nums []int, k int) int {
    var c int
    for i := range nums {
        var sum int
        for j := i; j >= 0 ; j-- {
            sum += nums[j]
            if sum == k {
                c++
            }
        }
    }
    return c
}

方法二：前缀和 + 哈希表优化

还有更快的算法么？

我们知道方法一的瓶颈在于对每个 i，我们需要枚举所有的 j 来判断是否符合条件。

除了通过加法累加 i 到 j 来判断 [j…i] 这个子数组和是否为 k，我们还可以通过前缀和的减法来判断。

我们定义 pre[i] 为 [0…i] 里所有数的和，则 pre[i] 可以由 pre[i−1] 递推而来，即：

pre[i]=pre[i−1]+nums[i]

那么「[j…i] 这个子数组和为 k 」这个条件我们可以转化为：

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

简单移项可得符合条件的下标 j 需要满足：

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

所以，当我们考虑以 i 结尾和为 k 的连续子数组个数时，只需要统计有多少个前缀和为 pre[i] - k （即 pre[j - 1]）的个数即可。

注意 j <= i，所以 pre[j-1] 表示 i 之前的前缀和。

具体做法如下：

• 使用 pre 变量记录前缀和，代表 pre[i]。
• 使用哈希表 hash 记录前缀和出现的次数。键值对为 pre[i] : pre_count。
• 从左到右遍历数组，计算当前前缀和 pre。
• 如果 pre - k 在哈希表中，则答案个数累加上 pre[pre - k]。
• 如果当前 pre 等于 k，则前缀和个数累加 1。
• 将当前前缀和 pre 记录到哈希表，即 hash[pre] += 1。
• 最后输出答案个数。

时间复杂度： O(n)，其中 n 为数组的长度。我们遍历数组的时间复杂度为 O(n)，中间利用哈希表查询删除的复杂度均为 O(1)，因此总时间复杂度为 O(n)。

空间复杂度： O(n)，其中 n 为数组的长度。哈希表在最坏情况下可能有 n 个不同的键值，因此需要 O(n) 的空间复杂度。

下面以 Golang 为例给出实现。

func subarraySum(nums []int, k int) int {
    m := make(map[int]int)
    var preSum int
    var c int
    for _, v := range nums {
        preSum += v
        c += m[preSum-k]
        if preSum == k {
            c++
        }
        // 将当前前缀和 pre 记录到哈希表。
        m[preSum]++
    }
    return c
}

---

参考文献

560. 和为K 的子数组