【哈希表】LeetCode 560. 和为K的子数组【中等】

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

 

示例 1：

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

示例 2：

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

提示：

1 <= nums.length <= 2 * 104
-1000 <= nums[i] <= 1000
-107 <= k <= 107

【分析】

方法一：暴力解法

该题目是要找到元素和等于k的所有连续子数组的个数。那么怎么求一个数组的连续子数组呢。可以使用两层循环来实现。

n = len(nums)

subnums = []
for i in range(n):
    for j in range(i, n):
        subnums.append(nums[i:j+1])

既然这样，稍作修改，就可以求元素和为k的连续子数组的个数了。

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int: 
        n = len(nums) # 求原始数组长度
        count = 0 # 初始化子数组个数

        for i in range(n):
            for j in range(i, n):
                if sum(nums[i: j+1]) == k:
                    count += 1        

        return count

提交代码，提示超出时间限制。

很明显，这么做的时间复杂度是O(n³)。因为除了两层for循环，还有一层求和。

其实，可以对上面方法做一些优化。因为我们求 nums[i: j + 1] 的下一次nums[i : j + 1 + 1]，其实只比前一次多了一个元素nums[j + 1]。那么我们就没必要每次都从i开始算了。但是这么做时间复杂度依然有O(n²)，依然超时。

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int: 
        n = len(nums) # 求原始数组长度
        count = 0 # 初始化子数组个数

        for i in range(n):
            sum = 0
            for j in range(i, n):
                sum += nums[j]
                if sum == k:
                    count += 1
        
        return count

前缀和

什么是前缀和：前缀和是指一个数组的某下标之前的所有数组元素的和（包含其自身）。

通常，会在前缀和首位放一个0，比如数组[1,2,3]，其前缀和是[0,1,3,6]。

前缀和通常可以帮助我们快速计算某个区间内的和。比如我们要计算i, j之间的和，那么就是nums[i] + nums[i+1] + ... + nums[j]，可以看作是:

nums[0] + nums[1] + ... + nums[i-1] + nums[i] + nums[i+1] + ... + nums[j] - 
nums[0] + nums[1] + ... + nums[i-1]

这个式子也是preSum[j] - preSum[i-1]。

那么，我们先遍历一次数组，求出前缀和数组。之后这个数组可以代替我们最开始暴力解法的sum函数。但是很遗憾，这种做法复杂度依然还是O(n²)。

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        # 要求的连续子数组
        count = 0
        n = len(nums)

        preSum = [0]

        # 求前缀和数组
        tmp = 0
        for i in range(n):
            tmp += nums[i]
            preSum.append(tmp)
        
        # 求和为k的连续子数组，求i到j之间的和
        for i in range(1, n+1):
            for j in range(i, n+1):
                if preSum[j] - preSum[i-1] == k:  # preSum[j] - preSum[i-1]代表着在nums数组中，前j个数之和减去前i-1个数之和
                    count += 1
        
        return count

进一步优化的话，我们可以边计算前缀和，边统计。遍历过程中，我们统计历史中每一个前缀和出现的次数，然后计算到位置i(含i)的前缀和preSum减去目标k在历史上出现过几次，假如出现过m次，代表第i位以前不含（不含i）有m个连续子数组的和为presum - k，这m个和为presum - k的连续子数组，每一个都可以和presum组合成为presum - （presum - k）= k

上述示意图，红色的是当前遍历到的前缀和presum，加入它之前有两个前缀和等于presume - k(蓝色区域范围)，那么很明显就会有两个连续子数组的元素和为k，对应橙色区域范围。

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        # 要求的连续子数组
        count = 0
        n = len(nums)
        preSums = collections.defaultdict(int)
        preSums[0] = 1

        presum = 0
        for i in range(n):
            presum += nums[i]

            # if preSums[presum - k] != 0
            count += preSums[presum - k] # 利用defaultdict的特性，当presum-k不存在时，返回的是0。这样避免了判断
            preSums[presum] += 1
        
        return count