问题描述
问题链接:https://leetcode.com/problems/subarray-sum-equals-k/
给定一个数组,请找到有多少个连续子数组加和等于给定目标数值?
leetcode,medium
560:Subarray Sum Equals K
问题分析
首先想到的是两件事情:
- 两层for循环遍历的解法;
- 两层for循环遍历大概率会出现TLE的错误,哈哈,不是废话嘛
因此,我们需要将O(n^2)的解法从时间上优化成O(n)的解法,这一般会增加空间上的开销,即“空间开销换时间开销”。这里参考的网上大神的常规做法,如下:
⚠️:下面提到的
累加和
,在不做特别说明的情况下,都是指从数组开头到当前位置处的所有元素的累加和。
- 申请字典数据结构
m
,初始时m中的元素为{0: 1}
(后面会解释原因),同时初始化满足条件的子数组的数量为res=0
; - 记录以当前位置编号为
i
的元素作为结尾的累加和sum
; - 从字典
m
中查询key=sum-k
对应的值m[sum-k]
,表示在位置编号为i
的元素之前已经出现了m[sum-k]
个元素作为结尾的累加和为sum-k
的情况。例如以位置编号为i-5
的元素结尾的累加和为sum-k
,那么从位置i-4
到i
之间的所有子数组的加和一定为k,刚好满足条件。此时以位置编号为i
的元素作为结尾的子数组(不一定从数组开头开始计算)加和等于k的情况一共是m[sum-k]
种,此时res
的更新只需要在原来的res基础上再加m[sum-k]
即可,即res = res + m[sum-k]
; - 对于累加和
sum
,其在字典结构m
中以sum
作为key的值应该自加1。
在步骤3中,当sum-k
刚好等于0时,说明此时从数组开头到当前位置所有元素累加和刚好等于k,满足条件,此时res
就要更新了。但是当第一次碰到这种情况时,res = res + m[sum-k]
中的 m[sum-k]=m[0]
=0(字典数据结构对于新的key的初始化其value为0)导致res无法更新,所以我们人为的在最开始时初始化{0: 1}
即可。
问题解法
直接看代码吧,结合问题分析
中的内容,代码应该很容易理解了。
def subarray_sum_equals_k(nums, target):
l = len(nums)
res = 0
m = {0: 1}
sum = 0
for i in range(l):
sum += nums[i]
res += m.get(sum - k, 0)
m[sum] = m.get(sum, 0) + 1
return res
测试代码如下:
if __name__ == "__main__":
nums = [1, 1, 1]
k = 2
print(subarray_sum_equals_k(nums, k))