和为 K 的子数组

标签: 数组 哈希表 前缀和

难度: Medium

给你一个整数数组 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

Submission

运行时间: 44 ms

内存: 18.7 MB

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        sum = 0
        sum_set = {0:1}
        res = 0

        for i in nums:
            sum += i
            res += sum_set.get(sum-k,0)

            if sum in sum_set:
                sum_set[sum] += 1
            else:
                sum_set[sum] = 1
        # print(sum_set)
        
        return res

Explain

这个题解使用了前缀和的思路。我们维护一个哈希表 sum_set,其中 key 为前缀和,value 为该前缀和出现的次数。我们遍历数组,对于当前遍历到的元素 i,我们将其加到 sum 中,sum 表示当前的前缀和。此时我们查询哈希表,看是否存在 key 为 sum-k 的元素,如果存在,说明从某个位置到当前位置的子数组和为 k,我们将 sum-k 出现的次数加到结果 res 中。最后我们将当前的前缀和 sum 加入哈希表中,如果哈希表中已经存在 sum,则将其出现次数加 1,否则将其加入哈希表,次数为 1。

时间复杂度: O(n)

空间复杂度: O(n)

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        # 初始化前缀和为 0,出现次数为 1
        sum = 0
        sum_set = {0:1}
        res = 0

        for i in nums:
            # 计算当前前缀和
            sum += i
            # 如果 sum-k 存在于哈希表中,说明存在和为 k 的子数组
            res += sum_set.get(sum-k,0)

            # 将当前前缀和加入哈希表
            if sum in sum_set:
                sum_set[sum] += 1
            else:
                sum_set[sum] = 1
        
        return res

Explore

前缀和是一个数组的所有元素的累加和,计算到当前元素为止。在这个算法中,通过维护一个前缀和来帮助快速计算任意子数组的和。具体来说,假设我们有一个前缀和 sum[i],它代表了从数组的开始到第 i 个元素的和。如果我们想计算从第 j 到第 i 的子数组的和(其中 j <= i),我们可以使用 sum[i] - sum[j-1] 来得到这个子数组的和。因此,通过查看已存储的前缀和,我们可以快速地计算出任意子数组的和,而无需重复遍历子数组。

在哈希表中设置键0的值为1是为了处理那些从数组开头到当前元素累加和正好为 k 的情况。此时,sum-k = 0,我们需要在哈希表中查找键为0的情况。如果不预先设置键0的值为1,那么在算法初始阶段,尽管可能存在一个有效的从数组起始到某点的和为k的子数组,我们也无法统计这种情况。因此,将0的值设置为1是为了正确处理累加和从开始就等于k的子数组。

在哈希表 sum_set 中,每个键代表一个前缀和,而值代表这个前缀和出现的次数。这样的设计允许算法快速查找在之前的数组元素中有多少次累加和等于当前 sum-k 的情况。每当这种情况出现时,这意味着从某个位置到当前位置的子数组之和为 k。通过记录每个前缀和出现的次数,我们可以快速计算出有多少子数组的和为 k。

算法处理包含负数或零的元素时的基本逻辑与处理全正数数组相同。前缀和的计算仍然是连续元素的累加,但前缀和的增减会受到元素值正负的影响。例如,负数会导致前缀和减少,而零不会改变前缀和。这种变化影响的是在哈希表中查找 sum-k 的机会。具体来说,负数可能导致更频繁的变化,从而增加找到匹配前缀和的机会,而零的存在则保持前缀和不变,可能在连续多个零的段中不改变子数组的数量。因此,不同的元素值(正数、负数、零)影响前缀和的变化,从而间接影响到统计和为k的子数组的数量。