[560][中等][前缀和][哈希] 和为K的子数组

题目描述

560. 和为K的子数组

给定一个整数数组和一个整数 k，你需要找到该数组中和为 k 的连续的子数组的个数。

示例 1 :

输入:nums = [1,1,1], k = 2
输出: 2 , [1,1] 与 [1,1] 为两种不同的情况。

说明 :

• 数组的长度为 [1, 20,000]。

• 数组中元素的范围是 [-1000, 1000] ，且整数 k 的范围是 [-1e7, 1e7]。

解题思路

暴力法是枚举出所有连续自数组, 判断是否符合条件, 遍历数组, 对于每个i, 都要枚举所有[j, i]这样的子数组. 优化就是通过前缀和的方法.

定义前缀和数组pre, pre[i]代表nums[0, 1, ..., i - 1]的和, 这样pre[i]可以由pre[i - 1]递推而来. 即pre[i] = pre[i - 1] + nums[i - 1]. pre前缀和数组的长度为n + 1, n是原数组nums的长度, pre[0]由于没有对应的前缀, 我们定义值为0.

那么每个连续子数组的和k就可以由前缀和数组得到: k = pre[i + 1] - pre[j], 代表[j, i]子数组之和. 因此在遍历到i位置时, 计算出前缀和pre[i + 1]的同时, 我们也可以计算出当前位置如果要满足子数组和k对应的pre[j], 而pre[j] = pre[i + 1] - j, 因此, 我们下一步的目标就是寻找值为pre[j] = pre[i + 1] - j的前缀和位置由多少.

因此, 为了搜索的快捷, 不会使用数组去存储前缀和, 而是使用哈希表. key值为前缀和值, value是当前前缀和对应的出现次数. 因此只需要遍历一次数组, 就能够得到满足条件的连续子数组的个数了.

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        mapping = {0: 1}
        ans, suffix = 0, 0
        for i, num in enumerate(nums):
            suffix += num
            resi = suffix - k
            if resi in mapping:
                ans += mapping[resi]
            mapping[suffix] = mapping[suffix] + 1 if suffix in mapping else 1
        return ans