利用前缀和处理a-b问题时的一个易错点

题目

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

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

示例 1:

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

示例 2:

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

提示:

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

题解

我们很容易可以想到使用一个前缀和数组presum[]处理,然后将区间内连续子列的和转化为前缀和的差,于是就转化为了A-B问题“A-B”问题的二分、哈希、双指针写法 ,利用哈希法解决。

代码如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class Solution {
public:
int subarraySum(vector<int>& nums, int k) {
int n=nums.size();
vector<int> presum(n+1,0);
presum[0]=0;
unordered_map<int,int> mp;
for(int i=1;i<=n;i++){
presum[i]=presum[i-1]+nums[i-1];
}
for(int x:presum){
mp[x]++;
}
int ans=0;
for(int x:presum){
ans+=mp[x-k];
}
return ans;
}
};

但是这里会遇到一个问题!当nums={1}, k=0时,会错误返回2。这是因为:将所有的前缀和一次性存入哈希表,包括当前处理的前缀和sum本身,这样会导致当sum -k等于sum时(比如k=0),会统计到自身,从而产生错误的计数。例如,当sum=0时,在第一次遍历presum数组时,哈希表中已经有了这个sum=0,当处理到它时,x-k=0,这时候会统计到至少一次,即使这个sum对应的位置可能没有更早的前缀和。

对代码进行修正:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class Solution {
public:
int subarraySum(vector<int>& nums, int k) {
int pre=0;
int ans=0;
unordered_map<int,int> cnt;
cnt[0]=1; //!!初始化前缀和为0的次数为1
for(int x:nums){
pre+=x;
ans+=cnt[pre-k];
cnt[pre]++;
}
return ans;
}
};

重点是避免自匹配:在计算ans+=cnt[pre-k]时,哈希表中仅包含遍历到当前元素之前的前缀和,确保不会将当前前缀和自身计入结果。


利用前缀和处理a-b问题时的一个易错点
http://example.com/2025/03/17/2025-03-17-利用前缀和处理a-b问题时的一个易错点/
作者
Kiriao
发布于
2025年3月17日
许可协议