Leetcode | 560. 和为 K 的子数组
560. 和为 K 的子数组
文章目录
- [560. 和为 K 的子数组](https://leetcode.cn/problems/subarray-sum-equals-k/)
-
- 题目
- 解法1：暴力枚举
- 解法2：前缀和
- 解法3：[官方题解](https://leetcode.cn/problems/subarray-sum-equals-k/solutions/238572/he-wei-kde-zi-shu-zu-by-leetcode-solution/)
欢迎关注公众号“三戒纪元”
题目
给你一个整数数组 nums 和一个整数 k ，请你统计并返回 该数组中和为 k 的连续子数组的个数 。
示例 1：
输入：nums = [1,1,1], k = 2
输出：2
示例 2：
输入：nums = [1,2,3], k = 3
输出：2
提示：
1 <= nums.length <= 2 * 104-1000 <= nums[i] <= 1000-107 <= k <= 107
解法1：暴力枚举
class Solution {
public:
int subarraySum(vector<int>& nums, int k) {
int count = nums.size();
if(count == 0) {
return 0;
}
int match = 0;
for(int i = 0; i < count; ++i) {
int sum = 0;
for(int j = i; j < count; ++j) {超时
sum += nums[j];
if(sum == k) {
++match;
}
}
}
return match;
}
};
但是这种解法对于计算包含N个1和N个0的数组会超时
即便是暴力穷举，官方题解思路也别出心裁，虽然同样会超时：
class Solution {
public:
int subarraySum(vector<int>& nums, int k) {
int count = 0;
for (int start = 0; start < nums.size(); ++start) {
int sum = 0;
for (int end = start; end >= 0; --end) {
sum += nums[end];
if (sum == k) {
count++;
}
}
}
return count;
}
};
作者：力扣官方题解
链接：https://leetcode.cn/problems/subarray-sum-equals-k/solutions/238572/he-wei-kde-zi-shu-zu-by-leetcode-solution/
来源：力扣（LeetCode）
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。
解法2：前缀和
class Solution {
public:
int subarraySum(vector<int>& nums, int k) {
int count = nums.size();
if(count == 0) {
return 0;
}
sums.resize(count + 1, 0);
int match = 0;
for(int i = 0; i < count; ++i) {
int sum = 0;
for(int j = i; j >= 0; --j) {
sum += nums[j];
if(sum == k) {
++match;
}
}
}
return match;
}
vector<int> sums;
};
同样会超时。
解法3：官方题解
我们可以基于方法一利用数据结构进行进一步的优化，我们知道方法一的瓶颈在于对每个 iii，我们需要枚举所有的 jjj 来判断是否符合条件，这一步是否可以优化呢？答案是可以的。
我们定义 p r e [ i ] pre[i] pre[i] 为 [ 0.. i ] [0..i] [0..i] 里所有数的和，则 p r e [ i ] pre[i] pre[i] 可以由 p r e [ i − 1 ] pre[i−1] pre[i−1] 递推而来，即：
p r e [ i ] = p r e [ i − 1 ] + n u m s [ i ] pre[i]=pre[i−1]+nums[i] pre[i]=pre[i−1]+nums[i]
那么「 [ j . . i ] [j..i] [j..i] 这个子数组和为 k 」这个条件我们可以转化为
p r e [ i ] − p r e [ j − 1 ] = = k pre[i]−pre[j−1]==k pre[i]−pre[j−1]==k
简单移项可得符合条件的下标 jjj 需要满足
p r e [ j − 1 ] = = p r e [ i ] − k pre[j−1]==pre[i]−k pre[j−1]==pre[i]−k
所以我们考虑以 i 结尾的和为 k 的连续子数组个数时只要统计有多少个前缀和为 p r e [ i ] − k pre[i]−k pre[i]−k 的 p r e [ j ] pre[j] pre[j]即可。我们建立哈希表 m p mp mp，以和为键，出现次数为对应的值，记录 p r e [ i ] pre[i] pre[i]出现的次数，从左往右边更新 m p mp mp 边计算答案，那么以 iii 结尾的答案 m p [ p r e [ i ] − k ] mp[pre[i]−k] mp[pre[i]−k]即可在 O ( 1 ) O(1) O(1) 时间内得到。最后的答案即为所有下标结尾的和为 k的子数组个数之和。
需要注意的是，从左往右边更新边计算的时候已经保证了 m p [ p r e [ i ] − k ] mp[pre[i]−k] mp[pre[i]−k]里记录的$ pre[j]$的下标范围是 0 ≤ j ≤ i 0≤j≤i 0≤j≤i 。同时，由于 p r e [ i ] pre[i] pre[i] 的计算只与前一项的答案有关，因此我们可以不用建立$ pre$ 数组，直接用 p r e pre pre变量来记录 p r e [ i − 1 ] pre[i−1] pre[i−1] 的答案即可。
class Solution {
public:
int subarraySum(vector<int>& nums, int k) {
int count = nums.size();
if(count == 0) {
return 0;
}
unordered_map<int, int> mp;
mp[0] = 1;
int match = 0;
int pre = 0;
for(int &num : nums) {
pre += num;
if(mp.find(pre - k) != mp.end() ) {
match += mp[pre - k];
}
mp[pre]++;
}
return match;
}
};