LeetCode 热题 100 | 560. 和为 K 的子数组

大家好，今天我们来解决一道经典的算法题——和为 K 的子数组。这道题在 LeetCode 上被标记为中等难度，要求我们统计数组中所有和为 k 的连续子数组的个数。下面我将详细讲解解题思路，并附上 Python 代码实现。

---

题目描述

给定一个整数数组 nums 和一个整数 k，统计并返回该数组中和为 k 的子数组的个数。子数组是数组中元素的连续非空序列。

示例：

输入：nums = [1,1,1], k = 2
输出：2
解释：子数组 [1,1] 和 [1,1] 的和为 2。

---

解题思路

我们需要统计数组中所有和为 k 的连续子数组的个数。直接暴力枚举所有子数组的时间复杂度为 O(n²)，对于较大的数组会超时。因此，我们需要一种更高效的方法。

核心思想：前缀和 + 哈希表

1. 前缀和：

   • 定义 prefix_sum 表示从数组开头到当前元素的累加和。
   • 对于任意子数组 nums[i..j]，其和可以表示为 prefix_sum[j] - prefix_sum[i-1]。

2. 哈希表优化：

   • 使用哈希表记录每个前缀和出现的次数。
   • 遍历数组时，计算当前前缀和 prefix_sum，并检查 prefix_sum - k 是否在哈希表中。如果存在，则说明存在若干个子数组的和为 k。

---

代码实现

from collections import defaultdict

def subarraySum(nums, k):
    # 初始化哈希表，记录前缀和出现的次数
    prefix_count = defaultdict(int)
    prefix_count[0] = 1  # 前缀和为 0 的情况出现一次

    prefix_sum = 0  # 当前前缀和
    count = 0  # 满足条件的子数组个数

    for num in nums:
        prefix_sum += num  # 更新前缀和
        # 如果 prefix_sum - k 在哈希表中，说明存在若干个子数组的和为 k
        if prefix_sum - k in prefix_count:
            count += prefix_count[prefix_sum - k]
        # 更新当前前缀和的出现次数
        prefix_count[prefix_sum] += 1

    return count

---

代码解析

1. 初始化哈希表：

   • prefix_count 用于记录每个前缀和出现的次数。
   • prefix_count[0] = 1 表示前缀和为 0 的情况出现一次（即空数组）。

2. 遍历数组：

   • 计算当前前缀和 prefix_sum。
   • 检查 prefix_sum - k 是否在哈希表中。如果存在，则说明存在若干个子数组的和为 k，累加到 count 中。
   • 更新当前前缀和 prefix_sum 的出现次数。

3. 返回结果：

   • 返回满足条件的子数组个数 count。

---

示例运行

示例 1

nums = [1, 1, 1]
k = 2
print(subarraySum(nums, k))  # 输出: 2

解释：

• 前缀和数组为 [1, 2, 3]。
• 满足条件的子数组为 [1, 1] 和 [1, 1]。

示例 2

nums = [1, 2, 3]
k = 3
print(subarraySum(nums, k))  # 输出: 2

解释：

• 前缀和数组为 [1, 3, 6]。
• 满足条件的子数组为 [1, 2] 和 [3]。

示例 3

nums = [1, -1, 1, -1]
k = 0
print(subarraySum(nums, k))  # 输出: 4

解释：

• 前缀和数组为 [1, 0, 1, 0]。
• 满足条件的子数组为 [1, -1], [-1, 1], [1, -1], 和 [1, -1, 1, -1]。

---

复杂度分析

• 时间复杂度：O(n)，其中 n 是数组的长度。我们只需要遍历数组一次。
• 空间复杂度：O(n)，哈希表最多存储 n 个前缀和。

---

总结

通过使用前缀和和哈希表，我们可以高效地统计数组中所有和为 k 的子数组的个数。这种方法的时间复杂度为 O(n)，能够处理较大的输入规模。希望这篇题解对你有帮助！如果还有其他问题，欢迎继续提问！

关注我，获取更多算法题解和编程技巧！