和为 K 的子数组

标签: 数组 哈希表 前缀和

难度: Medium

给定一个整数数组和一个整数 k请找到该数组中和为 k 的连续子数组的个数。

示例 1:

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

示例 2:

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

提示:

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

注意:本题与主站 560 题相同: https://leetcode-cn.com/problems/subarray-sum-equals-k/

Submission

运行时间: 48 ms

内存: 18.7 MB

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        cur_sum, count = 0, 0
        pre_dict = {0: 1}
        for i in range(len(nums)):
            cur_sum += nums[i]

            if cur_sum - k in pre_dict:
                count += pre_dict[cur_sum - k]

            if cur_sum in pre_dict:
                pre_dict[cur_sum] += 1
            else:
                pre_dict[cur_sum] = 1
        return count

Explain

本题解使用了前缀和与哈希表的组合策略。首先,定义一个变量 `cur_sum` 来存储从数组开始到当前元素的连续子数组的和。使用一个哈希表 `pre_dict` 来记录所有可能的前缀和的出现次数,其中键是前缀和,值是该前缀和出现的次数。在遍历数组的每个元素时,更新 `cur_sum`。接下来,检查 `cur_sum - k` 是否存在于 `pre_dict` 中,如果存在,表示找到了一个和为 `k` 的子数组,其次数为 `pre_dict[cur_sum - k]`。然后,更新哈希表,增加当前 `cur_sum` 的计数。这种方法能够有效地在一次遍历中计算出和为 `k` 的子数组的数量。

时间复杂度: O(n)

空间复杂度: O(n)

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        cur_sum, count = 0, 0 # 初始化当前和与计数器
        pre_dict = {0: 1} # 初始化哈希表,用于存储前缀和及其出现次数
        for i in range(len(nums)): # 遍历数组每个元素
            cur_sum += nums[i] # 更新当前前缀和

            if cur_sum - k in pre_dict: # 如果当前前缀和减去k在哈希表中存在
                count += pre_dict[cur_sum - k] # 累加对应的次数到计数器

            if cur_sum in pre_dict: # 如果当前前缀和已存在于哈希表
                pre_dict[cur_sum] += 1 # 更新该前缀和的计数
            else: # 如果当前前缀和不在哈希表中
                pre_dict[cur_sum] = 1 # 在哈希表中记录这个前缀和
        return count # 返回计数器的值

Explore

使用前缀和的主要原因是效率。如果直接遍历数组来计算每个可能的子数组的和,这种方法的时间复杂度会是O(n^2),因为对于数组中的每个元素,你可能都需要遍历后续的所有元素来计算子数组的和。相比之下,前缀和方法可以将这个问题简化为O(n)复杂度。通过使用前缀和,我们可以快速计算任何子数组的和,因为子数组 `sum(i, j)` 可以通过 `prefixSum[j] - prefixSum[i-1]` 得到,这使得算法更加高效。

在哈希表 `pre_dict` 中设置 `{0: 1}` 的意义在于,它代表了一个虚拟的前缀和,即在数组的开始之前,已经有一个和为0的前缀和出现了1次。这样做的主要目的是为了正确处理那些从数组第一个元素开始的子数组,其和恰好为k的情况。如果没有将0初始化在哈希表中,那么当遇到累积和恰好等于k的情况时,我们无法记录这是一个有效的子数组,因为没有0作为参照。

检查 `cur_sum - k` 是为了找出是否存在一个子数组的和正好等于k。这是基于前缀和的性质:如果两个前缀和之差等于k(即 `prefixSum[j] - prefixSum[i] = k`),那么从索引 `i+1` 到 `j` 的子数组和就是k。在这里,`cur_sum` 代表到当前元素为止的前缀和,通过查询 `cur_sum - k` 是否存在于哈希表中,我们实际上是在查找是否有一个之前的前缀和,使得从那个点到当前点的子数组和刚好为k。这是一种有效检测子数组和为k的方法,而且通过哈希表的帮助,这种检查可以非常迅速地完成。