0560. 和为 K 的子数组
https://leetcode-cn.com/problems/subarray-sum-equals-k/
给定一个整数数组和一个整数 k,你需要找到该数组中和为 k 的连续的子数组的个数。
示例 1 :
输入:nums = [1,1,1], k = 2
输出: 2 , [1,1] 与 [1,1] 为两种不同的情况。
说明 :
数组的长度为 [1, 20,000]。
数组中元素的范围是 [-1000, 1000] ,且整数 k 的范围是 [-1e7, 1e7]。
- 哈希表
- 前缀和
- 阿里
- 腾讯
- 字节
符合直觉的做法是暴力求解所有的子数组,然后分别计算和,如果等于 k,count 就+1.这种做法的时间复杂度为 O(n^2),代码如下:
class Solution:
def subarraySum(self, nums: List[int], k: int) -> int:
cnt, n = 0, len(nums)
for i in range(n):
sum = 0
for j in range(i, n):
sum += nums[j]
if (sum == k): cnt += 1
return cnt
实际上刚开始看到这题目的时候,我想“是否可以用滑动窗口解决?”。但是很快我就放弃了,因为看了下数组中项的取值范围有负数,这样我们扩张或者收缩窗口就比较复杂。第二个想法是前缀和,保存一个数组的前缀和,然后利用差分法得出任意区间段的和,这种想法是可行的,代码如下:
class Solution:
def subarraySum(self, nums: List[int], k: int) -> int:
cnt, n = 0, len(nums)
pre = [0] * (n + 1)
for i in range(1, n + 1):
pre[i] = pre[i - 1] + nums[i - 1]
for i in range(1, n + 1):
for j in range(i, n + 1):
if (pre[j] - pre[i - 1] == k): cnt += 1
return cnt
然而题目要求的仅仅是求总数,而不需要把所有的子数组求出。因此我们可直接使用下面的时间复杂度为 $O(n)$ 的算法。
这种做法的核心点在于, 题目让求的子数组总个数其实等价于:
- 以索引 0 结尾的子数组个数
- 以索引 1 结尾的子 数组个数
- 。。。
- 以索引 n - 1 结尾的子数组个数
而以索引 i 结尾的子数组个数等于:前缀和为 acc - k 的子数组个数,其中 acc 为当前的前缀和。为了能够快速取出前缀和为 acc - k 的个数,我们可将其存到哈希中。
具体算法: