前缀和 + 哈希表 O(n) 求解子数组和等于k的个数

题目描述

小午有一个长度为 $n$ 的数组 $a$，下标从 1 开始，小枫有一个整数 $k$。
小枫想知道数组 $a$ 的所有连续子序列中，有多少个元素之和等于 $k$。

形式化地，有多少对 $(l, r)$ 满足以下条件：

• $1 \le l \le r \le n$
• $\sum_{i=l}^{r} a_i = k$

输入格式

第一行输入两个整数 $n, k$，分别表示数组长度以及要求的元素之和。
第二行输入 $n$ 个整数 $a_i$，表示数组中第 $i$ 个元素的大小。

输出格式

输出一个整数，表示元素之和等于 $k$ 的连续子序列的个数。

解题思路

1. 前置知识：前缀和

定义前缀和数组 $s$：

• $s[0] = 0$
• $s[i] = a_1 + a_2 + \dots + a_i$（数组下标从 1 开始）

子数组 $a[l…r]$ 的和可以表示为：

$$\sum_{i=l}^{r} a_i = s[r] - s[l-1]$$

我们的目标是找到所有满足 $s[r] - s[l-1] = k$ 的 $(l, r)$ 对，等价于找有多少个 $s[l-1]$ 等于 $s[r] - k$。

2. 朴素暴力解法（超时）

• 预处理出完整的前缀和数组 $s$；
• 枚举所有子数组的起点 $l$ 和终点 $r$（$1 \le l \le r \le n$）；
• 计算 $s[r] - s[l-1]$ 并统计等于 $k$ 的次数。

缺点：时间复杂度为 $O(n^2)$，对于 $n=2e5$ 的数据范围会超时。

3. 前缀和 + 哈希表优化（最优解）

核心思想：遍历到位置 $r$ 时，只需知道前面有多少个 $s[l-1]$ 等于 $s[r] - k$，这个数量就是以 $r$ 为终点的、和为 $k$ 的子数组个数。

具体步骤：

1. 初始化哈希表 mp，mp[0] = 1（对应 $s[0]=0$，确保能统计从数组开头到当前位置的子数组）；
2. 遍历数组，实时计算当前前缀和 $s$（无需存储整个前缀和数组，节省空间）；
3. 查询哈希表中 $s - k$ 的出现次数，累加到答案 ans；
4. 将当前前缀和 $s$ 加入哈希表（计数 + 1）；
5. 遍历结束后，ans 即为最终结果。

代码实现

方法一：朴素暴力解法（O(n²) 超时）

#include<bits/stdc++.h>
using namespace std;
const int N = 2e5 + 9;

int n, a[N], ans;
long long k, s[N];

int main() {
    cin >> n >> k;
    for (int i = 1; i <= n; i++) {
        cin >> a[i];
        s[i] = s[i-1] + a[i];
    }
    
    for (int l = 1; l <= n; l++) {
        for (int r = l; r <= n; r++) {
            if (s[r] - s[l-1] == k) {
                ans++;
            }
        }
    }
    
    cout << ans << endl;
    return 0;
}

方法二：前缀和 + 哈希表优化 O（n） AC

#include<bits/stdc++.h>
using namespace std;
const int N = 2e5 + 9;

int n, a[N];
long long k, s, ans;
unordered_map<long long, int> mp;

int main() {
    cin >> n >> k;
    mp[0] = 1; // 初始化前缀和s[0]=0的计数为1
    
    for (int i = 1; i <= n; i++) {
        cin >> a[i];
        s += a[i]; // 实时计算当前前缀和
        
        // 查找有多少个前缀和等于s-k，累加到答案
        if (mp.count(s - k)) {
            ans += mp[s - k];
        }
        
        mp[s]++; // 将当前前缀和存入哈希表
    }
    
    cout << ans << endl;
    return 0;
}