给你一个整数数组 nums 和一个整数 k ,请你统计并返回 该数组中和为 k 的连续子数组的个数 。
示例 1:
输入:nums = [1,1,1], k = 2
输出:2
示例 2:
输入:nums = [1,2,3], k = 3
输出:2
提示:
1 <= nums.length <= 2 * 10^4
-1000 <= nums[i] <= 1000
-10^7 <= k <= 10^7
解法:前缀和 + 哈希表优化
参考官方题解:指路
思路:
考虑以i结尾和为k的连续子数组个数,我们需要统计符合条件的下标j的个数,其中0≤j≤i这个子数组的和恰好为k
图解(看懂图解,就自然思路很清晰了)
一开始,初始值:和为0、个数为1,所以mp:{{0,1}}
加入第一个元素3,也就是多了一对:和为3、累计和为3的个数是1,所以mp多了个(3,1),此时累计的和(pre)为3
加入第二个元素4,多了一对:和为7、累计和为7的个数是1,所以mp多了个(7,1),此时累计的和(pre)为7。由于pre-k为0,而0这个和在mp里出现过,就是初始值的那一对(0,1),其中0代表和,1代表个数。所以count+1。
加入第三个元素7,多了一对(14,1)。此时累计的和(pre)为14。由于pre-k为7,而7这个和在mp里出现过,就是第三对(7,1),其中7代表和,1代表个数。所以count再+1。
再加入第四个元素2,累计的和(pre)为16,pre-k=9,9没在mp里出现过,所以不用管。
再加入第五个元素-3,累计的和(pre)为13,pre-k=6,6没在mp里出现过,所以不用管。
重点!!加入第六个元素1,此时累计的和为14,14在mp里面已经存在啦,所以在原本的(14,1)变为(14,2)即可,不要再给mp新增一对(14,1)啦。此外!pre-k为7,7在mp里面出现过,所以count再+1。
加入第七个元素4。
加入第八个元素2。此时pre-k为13,13在mp里出现过,所以count再+1。最终count为4,结果就是4啦。
对应图例的代码实现:
class Solution:
def subarraySum(self, nums: List[int], k: int) -> int:
d = {}
acc = count = 0
for num in nums:
acc += num
if acc == k:
count += 1
if acc - k in d:
count += d[acc-k]
if acc in d:
d[acc] += 1
else:
d[acc] = 1
return count