LeetCode 560：和为 K 的子数组 | 前缀和与哈希表

LeetCode 560和为 K 的子数组 | 前缀和与哈希表引言和为 K 的子数组Subarray Sum Equals K是 LeetCode 第 560 题难度为 Medium。题目要求在给定整数数组中找出连续子数组的元素和等于 K 的数量。这道题是前缀和与哈希表结合的经典案例展示了如何将看似 O(n²) 的问题优化到 O(n) 时间复杂度。这道题的核心挑战在于如何高效地统计和等于 K 的子数组数量。如果使用暴力枚举所有子数组时间复杂度为 O(n²)。通过使用前缀和加哈希表的方法我们可以在 O(n) 时间内解决这个问题将时间复杂度降低一个数量级。问题分析题目描述给定一个整数数组 nums 和一个整数 K返回该数组中元素和等于 K 的连续子数组的数量。例如输入 nums [1, 1, 1]K 2输出为 2因为子数组 [1, 1]索引 0-1 和 1-2的和都等于 2。输入 nums [1, 2, 3]K 3输出为 2因为子数组 [1, 2] 和 [3] 的和都等于 3。问题特点这道题的关键在于连续子数组的要求。连续子数组意味着子数组的边界必须是连续的这使得前缀和技术成为可能。如果不要求连续问题就变成了计数问题需要用不同的方法处理。另一个挑战是如何在 O(n) 时间内统计所有和为 K 的子数组。暴力方法需要枚举所有可能的子数组起点和终点时间复杂度为 O(n²)。前缀和加哈希表的方法可以将时间复杂度降低到 O(n)。前缀和分析对于连续子数组 nums[i:j]包含 i 到 j其和等于 prefixSum[j1] - prefixSum[i]其中 prefixSum[k] 是前 k 个元素的和即 nums[0:k]。因此nums[i:j] 的和等于 K 当且仅当 prefixSum[j1] - prefixSum[i] K即 prefixSum[i] prefixSum[j1] - K。这个等式告诉我们对于位置 j1 的前缀和如果存在 i 使得 prefixSum[i] prefixSum[j1] - K那么以 i 为起点的子数组 nums[i:j] 的和就等于 K。因此我们只需要统计对于每个位置有多少个之前的前缀和等于当前前缀和减去 K。哈希表优化哈希表的设计我们使用哈希表来记录每个前缀和出现的次数。哈希表的键是前缀和的值值是该前缀和出现的次数。在遍历数组时对于每个位置我们首先计算当前的前缀和 current_sum。然后我们在哈希表中查找 key current_sum - K 的出现次数这个次数表示有多少个以当前位置结尾的和为 K 的子数组。def subarraySum(nums, k): prefix_sum 0 count 0 prefix_map {0: 1} for num in nums: prefix_sum num count prefix_map.get(prefix_sum - k, 0) prefix_map[prefix_sum] prefix_map.get(prefix_sum, 0) 1 return count初始化哈希表哈希表初始化为 {0: 1}表示前缀和为 0 出现了一次。这个初始化的作用是处理从数组开头开始的子数组。例如如果子数组 nums[0:j] 的和等于 K那么 prefixSum[j1] - prefixSum[0] K即 prefixSum[j1] K。由于我们初始化 prefixSum[0] 0在处理第一个元素时current_sum nums[0]我们需要检查 current_sum - K 是否等于 0即 K 是否等于 current_sum。如果没有初始化 {0: 1}这种情况就无法被正确计数。算法正确性证明数学归纳法我们使用数学归纳法证明算法的正确性。基例当处理第一个元素 nums[0] 时current_sum nums[0]。此时哈希表中只有 {0: 1}。如果 nums[0] K那么 prefix_sum - K 0 在哈希表中count 加上 1正确计数了子数组 nums[0:0]。如果 nums[0] ! Kcount 不变。基例成立。归纳假设假设在处理第 i 个元素之前哈希表正确记录了所有前缀和 prefixSum[0:i] 出现的次数count 正确计数了所有以第 i 个元素之前的位置结尾的和为 K 的子数组。归纳步骤处理第 i 个元素 nums[i] 后current_sum prefixSum[i1]。根据归纳假设此时 prefixSum[i] 已经在哈希表中。我们查找 prefix_sum - k prefixSum[i1] - K如果存在某个 j i1 使得 prefixSum[j] prefixSum[i1] - K那么子数组 nums[j:i1] 的和等于 prefixSum[i1] - prefixSum[j] K。因此将 prefix_map[prefix_sum - k] 加到 count 上正确计数了所有以 i 结尾的和为 K 的子数组。然后我们更新 prefix_map[current_sum]为后续的处理做准备。归纳步骤成立。复杂度分析时间复杂度时间复杂度为 O(n)因为我们只遍历数组一次每次迭代只进行常数次的哈希表操作查找和插入。空间复杂度空间复杂度为 O(n)在最坏情况下哈希表需要存储 n1 个不同的前缀和加上初始的 0。代码实现Python 实现def subarraySum(nums, k): prefix_sum 0 count 0 prefix_map {0: 1} for num in nums: prefix_sum num count prefix_map.get(prefix_sum - k, 0) prefix_map[prefix_sum] prefix_map.get(prefix_sum, 0) 1 return countJava 实现public int subarraySum(int[] nums, int k) { int prefixSum 0; int count 0; MapInteger, Integer prefixMap new HashMap(); prefixMap.put(0, 1); for (int num : nums) { prefixSum num; count prefixMap.getOrDefault(prefixSum - k, 0); prefixMap.put(prefixSum, prefixMap.getOrDefault(prefixSum, 0) 1); } return count; }边界情况处理空数组当数组为空时循环不会执行count 保持为 0返回 0符合预期。全负数数组算法同样适用于全负数数组。前缀和可以是负数哈希表的键可以是负数没有问题。全正数数组算法同样适用。对于全正数数组我们可以使用滑动窗口等其他方法但前缀和加哈希表的方法同样正确。K 为零当 K 0 时算法仍然有效。此时我们需要统计和为零的子数组数量。初始的 {0: 1} 会正确处理空前缀和。溢出问题在 Java 等语言中前缀和可能超出整数范围。由于 LeetCode 的测试用例通常在合理范围内且 Python 的整数可以自动扩展这个问题通常不需要特别处理。但在某些语言中可能需要使用 long 类型来存储前缀和。测试用例def test_subarray_sum(): assert subarraySum([1, 1, 1], 2) 2 assert subarraySum([1, 2, 3], 3) 2 assert subarraySum([1], 0) 0 assert subarraySum([0, 0], 0) 3 assert subarraySum([-1, -1, 1], 0) 1 assert subarraySum([1, 2, 3], 0) 0 print(所有测试用例通过)变种问题和至少为 K 的子数组如果问题改为统计和至少为 K 的子数组数量我们需要修改哈希表的查找逻辑。由于至少是一个范围查询不能简单地用哈希表解决。可以使用二分查找或滑动窗口等方法。乘积为 K 的子数组如果问题改为乘积为 K 的子数组需要使用对数变换将乘法转换为加法或者使用哈希表存储每个元素出现的次数适用于 K 包含较多因子的场景。和为 K 的最长子数组如果问题改为找到和为 K 的最长子数组的长度我们可以稍微修改算法在哈希表中存储每个前缀和首次出现的位置然后计算最长长度。实际应用场景和为 K 的子数组问题在现实中有很多应用。在金融领域可以用来检测连续交易日使累计收益达到特定值的情况。在信号处理中可以用来检测能量达到特定值的信号段。在数据分析中可以用来找出满足特定条件的数据段。前缀和加哈希表的方法是处理连续子数组统计问题的通用范式。掌握这个技巧后可以解决很多类似的问题。总结和为 K 的子数组问题展示了前缀和与哈希表结合的威力。通过将查找和为 K 的子数组转化为查找差值为 K 的两个前缀和我们可以在 O(n) 时间内解决看似需要 O(n²) 的问题。这个问题的关键洞察是子数组和等于 K 两个前缀和的差等于 K。通过哈希表快速查找满足条件的前缀和我们可以高效地统计所有符合条件的子数组。希望通过本文的讲解读者能够深入理解前缀和与哈希表结合的思想并将其应用到更多类似问题的解决中。