LeetCode10：和为K的数组

10.和为K的数组

题目

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

子数组是数组中元素的连续非空序列。

示例 1：

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

示例 2：

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

提示：

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

题解

  • 子数组是数组中连续的一段元素。

    例如，对于数组 nums = [a, b, c, d]，以下都是子数组：

    • [a]
    • [a, b]
    • [b, c, d]
    • [c]

    注意：子数组要求连续，因此 [a, c] 不是子数组，因为中间缺了 b。

一、暴力枚举

由于子数组的长度不确定，无法确定是由几个数字求得的K，采用暴力枚举，尝试所有可能。需要确定起点和终点：

外层循环确定子数组的起点：假设数组的长度为 n，我们可以让起点 i 从 0 到 n-1 遍历。

内层循环确定子数组的结束点：
对于固定起点 i，结束点 j 从 i 开始遍历到 n-1。这样可以确保子数组是连续的，即子数组为 nums[i] 到 nums[j]（包含两端）。

计算子数组的和：
对于每个确定的起点 i 和结束点 j，计算子数组 [nums[i], nums[i+1], ..., nums[j]] 的和，并检查是否等于 k。如果等于 k，则计数器加 1。

1
2
3
4
5
6
7
8
9
10
11
12
13
from typing import List
class Solution:
    # 暴力枚举
    def subarraySum(self, nums: List[int], k: int):
        length = len(nums)
        count = 0
        for i in range(length):
            cur_sum = 0
            for j in range(i,length):
                cur_sum += nums[j]  # 第一个和是0加上起点
                if cur_sum == k:
                    count += 1
        return count

二、前缀和+哈希表（优化）

子数组的和可以通过两个前缀和的差来计算。这是一种常见的算法技巧：如果我们维护每个索引的前缀和，那么对于每个索引，可以通过哈希表快速查找之前存在的前缀和（即当前前缀和减去k）。这种方法将时间复杂度从O(n^2)降低到O(n)。

  • 思路概述

    1. 前缀和
      定义 prefix[i] 为从数组开始到第 i 个元素（包含 i）的和。

      对于任意子数组 nums[i..j]，其和为：sum(i,j)=prefix[j]-prefix[i-1]。

      特别地，令 prefix[-1] = 0。

    2. 转换问题
      若要求 prefix[j] - prefix[i-1] = k，则可以转换为：prefix(i-1)=prefix[j]-k。

      也就是说，当遍历到第 j 个元素时，如果之前存在某个前缀和等于 prefix[j] - k，则从该位置之后到 j 的子数组之和为 k。

    3. 使用哈希表记录前缀和出现次数

      • 初始化哈希表（例如 Python 中的字典），令 prefix_count[0] = 1，表示空前缀和 0 出现一次。
      • 遍历数组，累计前缀和 prefix_sum，并检查 prefix_sum - k 是否在哈希表中出现过，若出现过，则将对应的次数累加到结果中。
      • 最后更新哈希表中当前前缀和 prefix_sum 的出现次数。

    这样可以将时间复杂度降低到 O(n)。

1
2
3
4
5
6
7
8
9
10
11
12
def subarraySum(self, nums: List[int], k: int):
        prefix_sum = 0
        prefix_count = {0:1} # 哈希表：键为前缀和，值为出现的次数，初始时 0 出现 1 次
        count = 0
        
        for num in nums:
            prefix_sum += num
            if prefix_sum - k in prefix_count:
                count += prefix_count[prefix_sum-k]
            prefix_count[prefix_sum] = prefix_count.get(prefix_sum,0) + 1
            
        return count

知识点

  • 子数组是否考虑去重？

    由题目给的示例，不需要。只要和为k都计入。

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

  • prefix_count字典的作用？

    用于记录遍历数组过程中，每个前缀和出现的次数。

    因为在遍历到索引j时，计算当前的前缀和prefix_sum，再从字典查找是否存在prefix_sum - k，字典中存储了该前缀和出现的次数，就可以直接累加到最终的计数中。

  • 初始化{0:1}的作用

    相当于一个虚拟的起点。因为在开始累加前，前缀和必定是 0。只有在累加过程中，当前缀和等于 k 时，通过 prefix_sum - k = 0 查找，才能正确计数那些从数组起始开始的满足条件的子数组。