给定一个数组,求和等于目标值的连续子数组的个数。
力扣中等题:560. 和为K的子数组
给定一个整数数组和一个整数 k,你需要找到该数组中和为 k 的连续的子数组的个数。
示例 1 :
输入:nums = [1,1,1], k = 2 输出: 2 , [1,1] 与 [1,1] 为两种不同的情况。
比较容易想到的是暴力解法,循环遍历得到所有的子数组的和,如果正好等于目标值则让计数加一,最后返回计数值。一般是用2个指针i
,j
分别指向指向子数组的开头和结尾,内外两层循环,时间复杂度为O(n^2)。
12345678910 | def subarraySum(self, nums: List[int], target: int) -> int: n=len(nums) res=0 for i in range(n): s=0 for j in range(i,n): s+=nums[j] if s==target: res+=1 return res |
---|
暴力法的时间复杂度过高,实际上我们可以预先处理数组,得到一个前缀和数组,在这个数组的基础上去计算,将时间复杂度降到O(n) 。
前缀和数组每一项对应的是数组从第0
项累加到第i
项的和,preSum[i]=sum(nums[0]+nums[1]+...+nums[i])
,比如[1,2,3,4,5]
的前缀和数组是[1,3,6,10,15]
。得到了这个数组,我们就可以以O(1)的时间复杂度求任意子数组的和,比如sum(nums[2],nums[3],nums[4])=preSum[4]-preSum[1]
,也就是说,对于任意的连续子数组nums[i,j]
,它的和等于preSum[j]-preSum[i-1]
。
我们可以为preSum
开头补充一项0
,这样preSum[i]
表示的意义为数组前i
个数字的和,连续子数组nums[i,j]
的和就可以表示为preSum[j+1]-preSum[i]
,省去了边界检查。
对于任意一个子数组的和为S[i,j]
等于目标值target
,有S[i,j]=preSum[j+1]-preSum[i]=target
,则preSum[i]=preSum[j+1]-target
。我们可以遍历preSum
数组,对于任意一个j
,记录对应有多少个preSum[i]
值可以满足条件preSum[i]=preSum[j+1]-target
,i<=j(i可以等于j,只有一项的数组我们也认为满足条件),总和即为结果。我们可以用一个哈希表来记录所有不同的preSum[i]
,同时存储个数,这样就省去了内循环的i
值遍历。
12345678910111213 | def subarraySum(self, nums: List[int], target: int) -> int: n=len(nums) preSum=[0]*(n+1) # 构造前缀和数组preSum for i in range(n): preSum[i+1]=nums[i]+preSum[i] counter=Counter() # 构造哈希表 counter[0]=1 # 初始化preSum[0]的个数为1 res=0 for j in range(n): if (k:=preSum[j+1]-target) in counter: # k相当于preSum[i] res+=counter[k] # 累加结果 counter[preSum[j+1]]+=1 # 当前的preSum值纳入统计 return res |
---|
可以看到,两个循环我们始终从preSum
的索引1
开始遍历,索引0
并没有用到,而且在第二个循环中,可以同步的去求preSum
的值,所以可以将两个循环省略为一个。
123456789101112 | def subarraySum(self, nums: List[int], target: int) -> int: n=len(nums) preSum=[0]*(n+1) # 构造前缀和数组preSum counter=Counter() # 构造哈希表 counter[0]=1 # 初始化preSum[0]的个数为1 res=0 for j in range(n): preSum[j+1]=nums[j]+preSum[j] if (k:=preSum[j+1]-target) in counter: # k相当于preSum[i] res+=counter[k] # 累加结果 counter[preSum[j+1]]+=1 # 当前的preSum值纳入统计 return res |
---|
在每次循环中,preSum
的当前值preSum[j+1]
是通过前一个值preSum[j]
计算出来的,也就是说每次循环中,我们只需要用到preSum
中的一个值即可,那就没必要存储整个preSum
数组,这也是常见的采用滚动数组压缩空间的方式。
123456789101112 | def subarraySum(self, nums: List[int], target: int) -> int: n=len(nums) preSum=0 # 压缩前缀和数组 counter=Counter() # 构造哈希表 counter[0]=1 # 初始化preSum[0]的个数为1 res=0 for j in range(n): preSum+=nums[j] if (k:=preSum-target) in counter: # k相当于preSum[i] res+=counter[k] # 累加结果 counter[preSum]+=1 # 当前的preSum值纳入统计 return res |
---|
最终,利用前缀和的思想和哈希表的数据结构,该题的时间复杂度为O(n),空间复杂度为哈希表的O(n)。
也可以用官方库函数计算前缀和数组preSum=list(itertools.accumulate(nums,initial=0))
等效于以下代码:
1234 | n=len(nums)preSum=[0]*(n+1) # 构造前缀和数组preSum for i in range(n): preSum[i+1]=nums[i]+preSum[i] |
---|