Leetcode.560 和为 K 的子数组
mid
给你一个整数数组 n u m s nums nums 和一个整数 k k k ,请你统计并返回 该数组中和为 k k k 的子数组的个数 。
子数组是数组中元素的连续非空序列。
输入:nums = [1,1,1], k = 2
输出:2
输入:nums = [1,2,3], k = 3
输出:2
我们假设 [ j , i ] [j,i] [j,i] 区间的子数组元素和为 k k k,即 :
n u m s [ j ] + n u m s [ j + 1 ] + . . . + n u m s [ i − 1 ] + n u m s [ i ] = k nums[j] + nums[j + 1] + ... + nums[i-1] + nums[i] = k nums[j]+nums[j+1]+...+nums[i−1]+nums[i]=k
我们用 s u m sum sum 表示 n u m s nums nums 的前缀和数组,可将上式转换为:
s u m [ i ] − s u m [ j − 1 ] = k sum[i] - sum[j-1] = k sum[i]−sum[j−1]=k
再转换一下得到:
s u m [ j − 1 ] = s u m [ i ] − k sum[j-1] = sum[i] - k sum[j−1]=sum[i]−k
那么以 n u m s [ i ] nums[i] nums[i] 为结尾的数组,我们只需要统计前面等于 s u m [ j − 1 ] sum[j-1] sum[j−1] 也就是 s u m [ i ] − k sum[i] - k sum[i]−k的前缀和的数量 t t t 即可。
那么这个 t t t 就是以 n u m s [ i ] nums[i] nums[i] 为结尾的数组中 和为 k k k 的子数组的数量。
我们只需要对每一个 n u m s [ i ] nums[i] nums[i] 都加上 t t t 即可,这样我们就可以统计出所有的 和为 k k k 的子数组的数量。
在实现上,我们使用哈希表来记录前缀和出现的次数。初始时,和为 0 0 0 ,也需要统计它的出现次数,即 { 0 , 1 } \{ 0 , 1 \} {0,1}。
时间复杂度: O ( n ) O(n) O(n)
C++代码:
class Solution {
public:
int subarraySum(vector<int>& nums, int k) {
int n = nums.size() , ans = 0 , sum = 0;
unordered_map<int,int> cnt;
cnt[0] = 1;
for(int i = 0;i < n;i++){
sum += nums[i];
ans += cnt[sum - k];
cnt[sum]++;
}
return ans;
}
};