Subarray Sum Equals K
给你一个整数数组 nums 和一个整数 k。请统计元素之和恰好等于 k 的子数组。子数组是由一个或多个相邻元素组成的一段。即使两个子数组包含相同的值,只要它们的起始位置或结束位置不同,就分别计数。元素的值可以是负数或 0。
函数
- numsinteger-array
- 整数数组,其中可以包含负数和零
- kinteger
- 子数组必须达到的总和才能计入
- 返回值integer
- 元素之和为 k 的子数组数量
约束条件
1 ≤ nums.length ≤ 2 × 104-1000 ≤ nums[i] ≤ 1000-107 ≤ k ≤ 107- 这么长的数组最多有 200,010,000 个子数组,因此答案可以用 32 位有符号整数表示。
示例
- 输入
- nums = [3, 4, -7, 1, 3, 3, 1, -4]k = 7
- 输出
- 4
- 解释
- 四次运行的结果相加为 7:
[3, 4]、[1, 3, 3]、[3, 3, 1]和[3, 4, -7, 1, 3, 3]。在最后一个结果中,-7 抵消了 3 和 4。
- 输入
- nums = [1, -1, 0]k = 0
- 输出
- 3
- 解释
- 有 3 个子数组的和为 0:
[1, -1]、[0]和整个数组[1, -1, 0]。子数组[-1, 0]的和为 -1,所以不计入。
- 输入
- nums = [2, 2, 2]k = 4
- 输出
- 2
- 解释
- 索引 0 和 1 处的连续段
[2, 2]与索引 1 和 2 处的连续段[2, 2]值相同,但位置不同,因此两者都要计入。
进阶
你会如何修改答案,使其仍在 O(n) 时间内返回总和为 k 的最长子数组的长度?
提示
逐个打开提示,每个都会多透露一点。
检查每个子数组是可行的,但 20,000 个数字大约有 2 亿个子数组。数值可能为负数,所以滑动窗口也行不通。你能用只需计算一次的数来表示任意子数组的和吗?
持续维护前缀和。两个位置之间元素的和,等于末尾位置的前缀和减去起始位置之前的前缀和。因此,以当前位置结尾的子数组之和恰好为
k,当且仅当之前的某个前缀和等于当前前缀和减去k。用一个哈希映射遍历数组一次,记录每个前缀和出现的次数,初始为空前缀:和为 0,出现 1 次。处理每个元素时,先将
prefix - k对应的计数加到答案中,然后再记录当前前缀和。
题解
检查最多 2 × 10^8 个子数组太慢了,而且负数也排除了使用滑动窗口的可能。子数组的和等于两个前缀和之差,因此,对于以当前位置结尾的子数组,每当较早的前缀和等于当前前缀和减去 k 时,就会有一个子数组的和等于 k。用哈希表可以在一次遍历中统计这些前缀和。
每次都从一个累计总数开始
结果正确,但无法在最大的测试中运行完毕
思路
固定第一个索引 start,然后向右移动 end,并将 nums[end] 加到累计的 total 中。total 始终保存从 start 到 end 的总和,因此遍历每一对索引就能检查每个子数组,每次只需进行一次加法。
即使总和超过 k,也不要停止。后面的负值可能会让总和降回来:在第一个示例中,从索引 0 开始的总和依次为 3、7、0、1、4、7,在索引 5 处第二次匹配。
开销取决于索引对的数量,当 n = 2 × 10^4 时,大约为 2 × 10^8。
算法
- 将
count设为 0。 - 对于从 0 到 n-1 的每个
start,将total设为 0。 - 对于从
start到 n-1 的每个end,将nums[end]加到total中。 - 如果
total等于k,就将count加 1,无论结果如何都继续。 - 返回
count。
def subarraySum(nums, k):
n = len(nums)
count = 0
for start in range(n):
total = 0
# Grow the subarray that begins at start, one element at a time
for end in range(start, n):
total += nums[end]
if total == k:
count += 1
return count使用计数映射的前缀和
思路
设 prefix[j] 表示前 j 个元素的和,对于空前缀,prefix[0] = 0。从索引 i 到索引 j-1 的子数组之和为 prefix[j] - prefix[i]。因此,每个等于当前前缀和减去 k 的较早前缀和,都标记着一个以当前位置结尾的匹配子数组的起点。要记录计数,而不是使用集合:每次重复出现的前缀和,都会开始一个不同的子数组。
在第一个示例中,前缀和为 0、3、7、0、1、4、7、8、4。第一个 7 找到一个 0,得到 [3, 4]。第二个 7 找到两个 0,得到 [3, 4, -7, 1, 3, 3] 和 [1, 3, 3]。8 找到一个 1,得到 [3, 3, 1]。
先查找,再记录当前前缀;否则当 k = 0 时,前缀会与自身匹配,形成一个空子数组。每个元素需要进行一次查找和一次更新:时间和空间复杂度均为 O(n)。
算法
- 创建一个映射
seen,其中seen[0] = 1,并将prefix和count设为 0。 - 对于每个元素,将它加到
prefix上。 - 将
seen[prefix - k]加到count上,缺少的键按 0 读取。 - 将
seen[prefix]加 1。 - 返回
count。
def subarraySum(nums, k):
# seen[p] = how many prefixes so far add up to p; the empty prefix adds up to 0
seen = {0: 1}
prefix = 0
count = 0
for num in nums:
prefix += num
# Each earlier prefix equal to prefix - k starts a subarray that ends here and sums to k
count += seen.get(prefix - k, 0)
seen[prefix] = seen.get(prefix, 0) + 1
return count
易错点和边界情况
- 当总和超过
k时缩小滑动窗口,但遇到负数就会出错:对于第一个示例,它返回 2 而不是 4。在暴力解法的内层循环中此时跳出,也会导致同样的错误。 - 漏掉
seen[0] = 1。对于nums = [5]和k = 5,它返回 0 而不是 1。 - 在查找之前记录当前前缀和。当
k = 0时,[1, -1, 0]返回 6 而不是 3。 - 用前缀和集合代替计数映射。对于
[0, 0, 0]和k = 0,它会漏掉重复项;答案是 6。
常见问题3
子数组和等于 K 的时间复杂度是多少?
前缀和与哈希映射的解法需要 O(n) 时间和 O(n) 额外空间:只需遍历一次,每个元素进行一次查找和一次更新。使用累加和检查每个子数组需要 O(n²) 时间。
为什么滑动窗口不适用于“子数组和等于 K”?
滑动窗口需要在窗口扩大时让总和增加,在窗口缩小时让总和减少。负数会破坏这一点:过大的窗口之后仍可能变成匹配项,因此没有规则能确定何时移动左边界。只有正数时,滑动窗口可以在 O(n) 时间和 O(1) 空间内解决这个问题。
为什么哈希表一开始会将 0 映射到 1?
这个条目表示第一个元素之前的空前缀,其和为 0。从索引 0 开始的子数组,其元素之和等于当前前缀和减去 0,因此如果没有这个条目,就永远不会统计这些子数组。对于 nums = [5] 和 k = 5,查找 5 - 5 = 0 时就能找到它。
相似题目
运用相同思路的题目。做两三道,才能真正掌握一种模式。
def subarraySum(nums, k):
# 在这里编写代码测试 #1
测试 #2
测试 #3
参数
nums = [3, 4, -7, 1, 3, 3, 1, -4] k = 7
预期输出
4