给你一个整数数组 nums 和一个整数 k ,请你统计并返回 该数组中和为 k 的子数组的个数 。
子数组是数组中元素的连续非空序列。
示例 1:
输入:nums = [1,1,1], k = 2
输出:2
示例 2:
输入:nums = [1,2,3], k = 3
输出:2
示例 3:
输入:nums = [-1,3,0,1], k = 2
输出:2
提示:
Medium。
★★★★☆
出题公司:阿里、腾讯、字节。
最容易想到是暴力枚举。
考虑以 i 结尾和为 k 的连续子数组个数,我们需要统计符合条件的下标 j 的个数,其中 0≤j≤i 且 [j…i] 这个子数组的和恰好为 k 。
可能有读者会认为假定我们确定了子数组的开头和结尾,还需要 O(n) 的时间复杂度遍历子数组来求和,那样复杂度就将达到 O(n^3) 从而无法通过所有测试用例。但是如果我们知道 [j,i] 子数组的和,就能 O(1) 推出 [j−1,i] 的和,因此这部分的遍历求和是不需要的,我们在枚举下标 j 的时候已经能 O(1) 求出 [j,i] 的子数组之和。
时间复杂度: O(n^2)
,其中 n 为数组的长度。枚举子数组开头和结尾需要 O(n^2)
的时间,其中求和需要 O(1) 的时间复杂度,因此总时间复杂度为 O(n^2)。
空间复杂度: O(1)。只需要常数空间存放若干变量。
下面以 Golang 为例给出实现。
func subarraySum(nums []int, k int) int {
var c int
for i := range nums {
var sum int
for j := i; j >= 0 ; j-- {
sum += nums[j]
if sum == k {
c++
}
}
}
return c
}
还有更快的算法么?
我们知道方法一的瓶颈在于对每个 i,我们需要枚举所有的 j 来判断是否符合条件。
除了通过加法累加 i 到 j 来判断 [j…i] 这个子数组和是否为 k,我们还可以通过前缀和的减法来判断。
我们定义 pre[i] 为 [0…i] 里所有数的和,则 pre[i] 可以由 pre[i−1] 递推而来,即:
pre[i]=pre[i−1]+nums[i]
那么「[j…i] 这个子数组和为 k 」这个条件我们可以转化为:
pre[i] − pre[j−1] == k
简单移项可得符合条件的下标 j 需要满足:
pre[j-1] == pre[i] - k
所以,当我们考虑以 i 结尾和为 k 的连续子数组个数时,只需要统计有多少个前缀和为 pre[i] - k (即 pre[j - 1])的个数即可。
注意 j <= i,所以 pre[j-1] 表示 i 之前的前缀和。
具体做法如下:
时间复杂度: O(n),其中 n 为数组的长度。我们遍历数组的时间复杂度为 O(n),中间利用哈希表查询删除的复杂度均为 O(1),因此总时间复杂度为 O(n)。
空间复杂度: O(n),其中 n 为数组的长度。哈希表在最坏情况下可能有 n 个不同的键值,因此需要 O(n) 的空间复杂度。
下面以 Golang 为例给出实现。
func subarraySum(nums []int, k int) int {
m := make(map[int]int)
var preSum int
var c int
for _, v := range nums {
preSum += v
c += m[preSum-k]
if preSum == k {
c++
}
// 将当前前缀和 pre 记录到哈希表。
m[preSum]++
}
return c
}