leetcode(560):和为K的子数组

leetcode(560):和为K的子数组

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

示例

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

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

前缀和+哈希表优化
我们定义pre[i]为[0…i]里所有数的和，则pre[i]可以由pre[i-1]递推而来，即
pre[i] = pre[i-1] + nums[i]
那么[j…i]这个子数组和为k 等价于
pre[i]-pre[j-1] == k
pre[j-1] == pre[i]-k
所以我们考虑以 i 结尾的和为 kk的连续子数组个数时只要统计有多少个前缀和为 pre[i]-k 的 pre[j] 即可。

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        res = sumi = 0
        n = len(nums)
        d = {0:1}
        for i in range(n):
            sumi += nums[i]
            sumj = sumi-k
            if sumj in d:
                res += d[sumj]
            d[sumi] = d.get(sumi,0)+1
        return res