【前缀和】560. 和为 K 的子数组

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

 

示例 1：

输入：nums = [1,1,1], k = 2
输出：2
示例 2：

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

---

 

思路：前缀和

前缀和定义：定义前缀和数组pre， 使得pre[i] 为 nums[0],  ...  , nums[i]  的所有数的和，  pre[i]  可以从pre[i-1]  得到， 即 pre[i] = pre[i-1] + nums[i]

 

 

from collections import defaultdict
class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        pre = defaultdict(int)   # 记录前缀和出现的次数
        pre[0] = 1  # 前缀和为0的出现1次
        ans = 0
        cur = 0    # 记录前缀和   
        for num in nums:
            cur += num
            if cur - k in pre:   # 由于使用的是 defaultdict(int)，字典中不存在key时，默认值为0，此if 可以省略。
                ans += pre[cur-k]
            pre[cur] += 1
        return ans