Sign in Subscribe
算法

和为 K 的子数组

loser

• 6 min read
Share

给你一个整数数组 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

前缀和-哈希表：

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        unordered_map<int, int> cnt;
        cnt[0] = 1;

        int pre = 0;
        int ans = 0;

        for (int num : nums)
        {
            pre += num;

            if (cnt.find(pre - k) != cnt.end())
            {
                ans += cnt[pre - k];
            }

            cnt[pre]++;
        }

        return ans;
    }
};

核心思想

这题要求的是 连续子数组 的和。

一看到连续区间求和，就可以优先考虑前缀和。

设：

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

那么从下标 left 到下标 right 的子数组和就是：

nums[left] + nums[left + 1] + ... + nums[right]

可以用前缀和表示为：

pre[right] - pre[left - 1]

题目要求这个值等于 k，所以：

pre[right] - pre[left - 1] = k

移项得到：

pre[left - 1] = pre[right] - k

也就是说，当我们遍历到当前位置 right，并且知道当前前缀和是 pre[right] 时，只要前面出现过前缀和 pre[right] - k，就说明存在一个以 right 结尾、和为 k 的子数组。

所以问题就变成了：

遍历数组时，统计之前有多少个前缀和等于 当前前缀和 - k。

这正好可以用哈希表解决。

为什么不能直接用滑动窗口

有些连续子数组问题可以使用滑动窗口，但这题不适合直接使用普通滑动窗口。

原因是：nums[i] 可以是负数。

如果数组里全是正数，那么窗口变大，窗口和一定变大；窗口变小，窗口和一定变小。

这样就可以通过移动左右指针控制窗口和。

但如果数组中有负数，这个单调性就不存在了：

  • 右边加入一个负数，窗口和可能变小
  • 左边移除一个负数，窗口和可能变大

所以不能简单地根据当前和大于或小于 k 来移动窗口。

因此，这题更稳定的做法是前缀和 + 哈希表。

哈希表存什么

代码中：

unordered_map<int, int> cnt;

表示：

  • 键：某个前缀和
  • 值：这个前缀和在之前出现过多少次

例如：

cnt[x] = 3;

表示在当前下标之前，前缀和 x 出现过 3 次。

当当前前缀和为 pre 时，我们需要找：

pre - k

如果 pre - k 在之前出现过 cnt[pre - k] 次，那么就说明有 cnt[pre - k] 个不同的起点，可以和当前下标组成和为 k 的子数组。

所以：

ans += cnt[pre - k];

然后再把当前前缀和加入哈希表：

cnt[pre]++;

为什么要初始化 cnt[0] = 1

这是这题最容易漏掉的一点。

cnt[0] = 1 表示：

在数组开始之前，存在一个前缀和为 0 的位置。

这样可以处理从下标 0 开始的子数组。

例如：

nums = [1, 2, 3], k = 3

当遍历到下标 1 时：

pre = 1 + 2 = 3

此时整个子数组 [1, 2] 的和就是 3。

根据公式，我们需要找：

pre - k = 3 - 3 = 0

如果没有提前设置 cnt[0] = 1，这个从下标 0 开始的合法子数组就会被漏掉。

所以初始化非常必要。

公式推导

为了让下标 left ... right 的子数组和等于 k，需要满足：

sum(left, right) = k

根据前缀和定义：

sum(left, right) = pre[right] - pre[left - 1]

所以：

pre[right] - pre[left - 1] = k

移项：

pre[left - 1] = pre[right] - k

当我们固定右端点 right 时，pre[right] 是确定的。

因此，合法子数组的个数就等于：

在 right 之前，有多少个前缀和等于 pre[right] - k。

这就是代码中：

ans += cnt[pre - k];

的来源。

正确性证明

我们证明：算法返回的 ans 等于数组中所有和为 k 的子数组个数。

结论 1：算法统计到的每个子数组都合法

当算法遍历到某个位置 right 时，当前前缀和是 pre。

如果哈希表中存在 pre - k，说明在当前下标之前，存在某个位置的前缀和为：

pre - k

设这个位置对应的是 left - 1。

那么子数组 nums[left ... right] 的和就是：

pre - (pre - k) = k

所以每次通过 cnt[pre - k] 增加的答案，都对应一个合法的和为 k 的子数组。

结论 2：每个合法子数组都会被算法统计到

任意一个和为 k 的子数组，设它的范围是：

nums[left ... right]

那么根据前缀和公式：

pre[right] - pre[left - 1] = k

也就是：

pre[left - 1] = pre[right] - k

当算法遍历到 right 时，left - 1 一定已经在之前被处理过。

因此哈希表中一定已经记录了这个前缀和 pre[left - 1]。

算法会通过查询：

cnt[pre[right] - k]

把这个子数组统计进答案。

所以所有合法子数组都不会被漏掉。

结论 3：重复前缀和需要计数

同一个前缀和可能出现多次。

例如，不同位置的前缀和都等于 x。

当当前前缀和是 x + k 时，这些位置都可以作为不同的子数组起点。

所以哈希表不能只记录某个前缀和是否出现过，而要记录它出现了多少次。

这也是为什么使用：

unordered_map<int, int> cnt;

而不是：

unordered_set<int> st;

得出结论

由结论 1 可知，算法不会统计错误子数组。

由结论 2 可知，算法不会漏掉合法子数组。

由结论 3 可知，算法可以正确处理多个相同前缀和带来的多个不同子数组。

因此算法正确。

举例理解

以：

nums = [1, 1, 1], k = 2

为例。

初始：

cnt[0] = 1

遍历第一个 1：

  • pre = 1
  • 查找 pre - k = -1，没有
  • 记录 cnt[1]++

遍历第二个 1：

  • pre = 2
  • 查找 pre - k = 0
  • cnt[0] = 1，说明找到子数组 [1, 1]
  • ans = 1
  • 记录 cnt[2]++

遍历第三个 1：

  • pre = 3
  • 查找 pre - k = 1
  • cnt[1] = 1，说明找到另一个子数组 [1, 1]
  • ans = 2
  • 记录 cnt[3]++

最终答案是 2。

复杂度分析

数组中的每个元素只会被遍历一次。

每次遍历时，哈希表的查询和更新平均时间复杂度都是 O(1)。

所以总时间复杂度是：

O(n)

哈希表最多记录 n + 1 个前缀和，所以空间复杂度是：

O(n)