题目链接 · 灵神原题解（署名来源）

前置题目

本题是连续子数组和问题，可以用前缀和处理。

请先完成前缀和模板题 303. 区域和检索 - 数组不可变，并阅读 我的题解。

转化

回顾 303 题的前缀和的定义：$s[0]=0,s[i]=\mathit{\text{nums}}[0]+\mathit{\text{nums}}[1]+\cdots +\mathit{\text{nums}}[i-1]$。

注意 $s$ 是一个长为 $n+1$ 的数组，第一个数是 $0$。

设 $i<j$，如果 $\mathit{\text{nums}}[i]$ 到 $\mathit{\text{nums}}[j-1]$ 的元素和等于 $k$，用前缀和表示，就是

$$s[j]-s[i]=k$$

问题转化为：

• $s$ 中有多少对下标 $(i,j)$ 满足 $0\le i<j\le n$ 且 $s[j]-s[i]=k$？

写成 $s[j]+(-s[i])=k$ 就能看得更明白，这是梦开始的地方——1. 两数之和。不过那题只需找到一对下标，而本题需要计算所有满足条件的下标对的个数。

枚举右，维护左

以 $\mathit{\text{nums}}=[1,1,-1,1,-1]$，$k=1$ 为例，其前缀和数组为 $s=[0,1,2,1,2,1]$。画出前缀和数组的折线图，如下：

如果用二重循环暴力枚举有多少个 $s[j]-s[i]=k$，时间复杂度是 $\mathcal{O}(n^2)$，太慢了。如何加速？

从两数之和中，我们可以学到什么？我们可以把 $s[j]-s[i]=k$ 移项，得

$$s[i]=s[j]-k$$

枚举当前的前缀和 $s[j]$，看看曾经有多少个前缀和等于 $s[j]-k$（配对）。每当我们在左边找到一个值等于 $s[j]-k$ 的前缀和，就找到了一个和为 $k$ 的子数组（因为 $s[j]-(s[j]-k)=k$）。

比如 $s[j]=2$，那么 $s[i]=s[j]-k=2-1=1$，我们要找的是 $j$ 左边有多少个 $s[i]=1$。在上面的例子中，遍历到 $s[4]=2$ 时，我们知道左边有 $2$ 个 $s[i]=1$，所以新找到了 $2$ 个和为 $1$ 的子数组。

用这个视角，再来算算上图中的那 $6$ 个和为 $1$ 的子数组。

$j$	$s[j]$	$s[j]-k$	$s[j]-k$ 的个数	解释
$0$	$0$	$-1$	$0$	无
$1$	$1$	$0$	$1$	$s[0]=0$
$2$	$2$	$1$	$1$	$s[1]=1$
$3$	$1$	$0$	$1$	$s[0]=0$
$4$	$2$	$1$	$2$	$s[1]=s[3]=1$
$5$	$1$	$0$	$1$	$s[0]=0$

一共有 $0+1+1+1+2+1=6$ 个和为 $k=1$ 的子数组。

具体来说，在遍历 $s[j]$ 的同时，用一个哈希表 $\mathit{\text{cnt}}$ 统计 $s[j]$ 的个数。哈希表的 key 是 $s[j]$，value 是值为 $s[j]$ 的前缀和的个数。遍历到 $s[j]$ 时，从哈希表中就可以找到有 $\mathit{\text{cnt}}[s[j]-k]$ 个 $s[i]$，加入答案。请读者动手算算上面的例子，加深理解。

答疑

问：为什么这样做可以不重不漏地计算？

答：暴力做法是，外层循环枚举 $j$，内层循环枚举 $i$，如果 $s[j]-s[i]=k$，那么答案加一。我们保留了「外层循环枚举 $j$」这个过程，把内层循环用哈希表优化成了 $\mathcal{O}(1)$，所以本质是对暴力算法的哈希表优化。既然暴力算法是不重不漏地计算，那么优化做法也是不重不漏地计算。

问：为什么要把 $s[0]=0$ 也加到哈希表中？

答：举个最简单的例子，$\mathit{\text{nums}}=[1],k=1$。如果不把 $s[0]=0$ 加到哈希表中，按照我们的算法，没法算出这里有 $1$ 个符合要求的子数组。也可以这样理解，要想把任意子数组都表示成两个前缀和的差，必须添加 $s[0]=0$，否则当子数组是前缀时，没法减去一个数，具体见 前缀和及其扩展 中的讲解。

问：为什么代码中要先更新 $\mathit{\text{ans}}$，再更新 $\mathit{\text{cnt}}$？这两行代码能否交换？

答：不行，这会在 $k=0$ 的时候算错。例如 $\mathit{\text{nums}}=[2],k=0$，正确答案应该是 $0$，但如果先把 $\mathit{\text{cnt}}[2]$ 加一，再把 $\mathit{\text{cnt}}[2]$ 加到 $\mathit{\text{ans}}$ 中，最后返回的 $\mathit{\text{ans}}$ 就不是 $0$ 了。

问：为什么这题不适合用滑动窗口做？

答：滑动窗口需要满足单调性，当右端点元素进入窗口时，窗口元素和是不能减少的。本题 $\mathit{\text{nums}}$ 包含负数，当负数进入窗口时，窗口左端点该往哪个方向移动？无法确定。如果没有负数的话，则可以用滑动窗口（恰好型滑动窗口），见 930. 和相同的二元子数组。

写法一：两次遍历

Python 解法1

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        s = [0] * (len(nums) + 1)
        for i, x in enumerate(nums):
            s[i + 1] = s[i] + x

        cnt = defaultdict(int)
        ans = 0
        for sj in s:
            ans += cnt[sj - k]
            cnt[sj] += 1
        return ans

C++ 解法1

// C++ 版待补充

C++ 解法2

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        int n = nums.size();
        vector<int> s(n + 1);
        for (int i = 0; i < n; i++) {
            s[i + 1] = s[i] + nums[i];
        }

        unordered_map<int, int> cnt;
        int ans = 0;
        for (int sj : s) {
            // 注意不要直接 += cnt[sj-k]，如果 sj-k 不存在，会插入 sj-k
            ans += cnt.contains(sj - k) ? cnt[sj - k] : 0;
            cnt[sj]++;
        }
        return ans;
    }
};

复杂度分析

• 时间复杂度：$\mathcal{O}(n)$，其中 $n$ 为 $\mathit{\text{nums}}$ 的长度。
• 空间复杂度：$\mathcal{O}(n)$。

写法二：一次遍历 · 其一

我们可以一边计算前缀和，一边遍历前缀和。

在遍历 $\mathit{\text{nums}}$ 之前，我们需要先统计 $s[0]=0$，即空前缀的元素和等于 $0$。往 $\mathit{\text{cnt}}$ 中添加 $\mathit{\text{cnt}}[0]=1$。

对比一下，两次遍历的代码循环了 $n+1$ 次，下面的代码循环了 $n$ 次，少的那一次是什么？就是对 $s[0]=0$ 的统计。

Python 解法3

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        cnt = defaultdict(int)
        cnt[0] = 1  # s[0]=0 单独统计
        ans = s = 0
        for x in nums:
            s += x
            ans += cnt[s - k]
            cnt[s] += 1
        return ans

C++ 解法3

// C++ 版待补充

C++ 解法4

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        unordered_map<int, int> cnt = {{0, 1}}; // s[0]=0 单独统计
        int ans = 0, s = 0;
        for (int x : nums) {
            s += x;
            // 注意不要直接 += cnt[s-k]，如果 s-k 不存在，这会插入 s-k，消耗更多空间
            ans += cnt.contains(s - k) ? cnt[s - k] : 0;
            cnt[s]++;
        }
        return ans;
    }
};

写法三：一次遍历 · 其二

在同一轮循环中，先把 $s[i-1]$ 加入哈希表，再根据 $s[i]$ 更新答案。

这样写无需初始化 $\mathit{\text{cnt}}[0]=1$。

Python 解法5

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        cnt = defaultdict(int)
        ans = s = 0
        for x in nums:
            cnt[s] += 1
            s += x
            ans += cnt[s - k]
        return ans

C++ 解法5

// C++ 版待补充

C++ 解法6

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        unordered_map<int, int> cnt;
        int ans = 0, s = 0;
        for (int x : nums) {
            cnt[s]++;
            s += x;
            // 注意不要直接 += cnt[s-k]，如果 s-k 不存在，这会插入 s-k，消耗更多空间
            ans += cnt.contains(s - k) ? cnt[s - k] : 0;
        }
        return ans;
    }
};

复杂度分析

• 时间复杂度：$\mathcal{O}(n)$，其中 $n$ 为 $\mathit{\text{nums}}$ 的长度。
• 空间复杂度：$\mathcal{O}(m)$，其中 $m$ 为不同前缀和的个数。如果设置了哈希表的容量，则空间复杂度为 $\mathcal{O}(n)$。

变形题

1. 改成计算元素和等于 $k$ 的最短子数组，要怎么做？
2. 改成计算元素和等于 $k$ 的最长子数组，要怎么做？
3. 改成计算元素和等于 $k$ 的所有子数组的长度之和，要怎么做？
4. 改成元素和至多为 $k$，要怎么做？见 363. 矩形区域不超过 K 的最大数值和。
5. 改成计算元素和为奇数的子数组个数，要怎么做？

欢迎在评论区分享你的思路/代码。

提示：思考题 4 可以枚举上下边界，转成一维数组。

分类题单

如何科学刷题？

1. 滑动窗口与双指针（定长/不定长/单序列/双序列/三指针/分组循环）
2. 二分算法（二分答案/最小化最大值/最大化最小值/第K小）
3. 单调栈（基础/矩形面积/贡献法/最小字典序）
4. 网格图（DFS/BFS/综合应用）
5. 位运算（基础/性质/拆位/试填/恒等式/思维）
6. 图论算法（DFS/BFS/拓扑排序/基环树/最短路/最小生成树/网络流）
7. 动态规划（入门/背包/划分/状态机/区间/状压/数位/数据结构优化/树形/博弈/概率期望）
8. 常用数据结构（前缀和/差分/栈/队列/堆/字典树/并查集/树状数组/线段树）
9. 数学算法（数论/组合/概率期望/博弈/计算几何/随机算法）
10. 贪心与思维（基本贪心策略/反悔/区间/字典序/数学/思维/脑筋急转弯/构造）
11. 链表、二叉树与回溯（前后指针/快慢指针/DFS/BFS/直径/LCA/一般树）
12. 字符串（KMP/Z函数/Manacher/字符串哈希/AC自动机/后缀数组/子序列自动机）

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府