前缀和 + 哈希表 O(N) 求解子数组和等于K的个数
题目描述
小午有一个长度为 nnn 的数组 aaa,下标从 1 开始,小枫有一个整数 kkk。
小枫想知道数组 aaa 的所有连续子序列中,有多少个元素之和等于 kkk。
形式化地,有多少对 (l,r)(l, r)(l,r) 满足以下条件:
* 1≤l≤r≤n1 \le l \le r \le n1≤l≤r≤n
* ∑i=lrai=k\sum_{i=l}^{r} a_i = k∑i=lr ai =k
输入格式
第一行输入两个整数 n,kn, kn,k,分别表示数组长度以及要求的元素之和。
第二行输入 nnn 个整数 aia_iai ,表示数组中第 iii 个元素的大小。
输出格式
输出一个整数,表示元素之和等于 kkk 的连续子序列的个数。
解题思路
1. 前置知识:前缀和
定义前缀和数组 sss:
* s[0]=0s[0] = 0s[0]=0
* s[i]=a1+a2+⋯+ais[i] = a_1 + a_2 + \dots + a_is[i]=a1 +a2 +⋯+ai (数组下标从 1 开始)
子数组 a[l…r]a[l…r]a[l…r] 的和可以表示为:
∑i=lrai=s[r]−s[l−1]\sum_{i=l}^{r} a_i = s[r] - s[l-1] i=l∑r ai =s[r]−s[l−1]
我们的目标是找到所有满足 s[r]−s[l−1]=ks[r] - s[l-1] = ks[r]−s[l−1]=k 的 (l,r)(l, r)(l,r) 对,等价于找有多少个 s[l−1]s[l-1]s[l−1] 等于 s[r]−ks[r] - ks[r]−k。
2. 朴素暴力解法(超时)
* 预处理出完整的前缀和数组 sss;
* 枚举所有子数组的起点 lll 和终点 rrr(1≤l≤r≤n1 \le l \le r \le n1≤l≤r≤n);
* 计算 s[r]−s[l−1]s[r] - s[l-1]s[r]−s[l−1] 并统计等于 kkk 的次数。
缺点:时间复杂度为 O(n2)O(n^2)O(n2),对于 n=2e5n=2e5n=2e5 的数据范围会超时。
3. 前缀和 + 哈希表优化(最优解)
核心思想:遍历到位置 rrr 时,只需知道前面有多少个 s[l−1]s[l-1]s[l−1] 等于 s[r]−ks[r] - ks[r]−k,这个数量就是以 rrr 为终点的、和为 kkk 的子数组个数。
具体步骤:
1. 初始化哈希表 mp,mp[0] = 1(对应 s[0]=0s[0]=0s[0]=0,确保能统计从数组开头到当前位置的子数组);
2. 遍历数组,实时计算当前前缀和 sss(无需存储整个前缀和数组,节省空间);
3. 查询哈希表中 s−ks - ks−k 的出现次数,累加到答案 ans;
4. 将当前前缀和 sss 加入哈希表(计数 + 1);
5. 遍历结束后,ans 即为最终结果。
代码实现
方法一:朴素暴力解法(O(N²) 超时)
###方法二:前缀和 + 哈希表优化 O(n) AC