题目来自LeetCode,链接:subarray-sum-equals-k。具体描述为:给定一个整数数组和一个整数 k,你需要找到该数组中和为 k 的连续的子数组的个数。
示例:
输入:nums = [1,1,1], k = 2
输出: 2 , [1,1] 与 [1,1] 为两种不同的情况。
一个很容易想到的思路就是暴力法,直接遍历所有可能的子数组,最好也只能做到时间复杂度为 O ( N 2 ) O(N^{2}) O(N2),java版可以AC但python就不行了。
另一种方法就是考虑数组各个位置上的累加和,比如[1,1,1]
的累加和就是[1,2,3]
,可以看到,如果一个连续子数组(假设从i到j)的和为k,那么在累加和数组(假设为S)中,必定有S[j]-S[i]=k
,上面的例子就有3-1=2
以及2-0=2
(0即子数组从头开始的情况)。因此可以在遍历累加和数组的过程中用一个字典/Map保存当前累加和a需要的下一个累加和a+k以使得两者之差为k,具体就是以当前累加和需要的下一个累加和a+k为key,需要a+k的次数为value。遍历时只要发现当前累加和在字典/Map中,对应的value值就是符合条件的子数组数,累加起来即可。此方法的时间复杂度只有
O
(
N
)
O(N)
O(N)。
JAVA版代码如下:
class Solution {
public int subarraySum(int[] nums, int k) {
int numsLen = nums.length, count = 0;
int[] cumulativeSum = new int[numsLen];
cumulativeSum[0] = nums[0];
for (int i = 1; i < numsLen; ++i) {
cumulativeSum[i] = cumulativeSum[i - 1] + nums[i];
}
Map<Integer, Integer> cumuSum2count = new HashMap<>();
cumuSum2count.put(k, 1);
for (int i : cumulativeSum) {
if (cumuSum2count.containsKey(i)) {
count += cumuSum2count.get(i);
}
int nextCumuSum = i + k;
if (cumuSum2count.containsKey(nextCumuSum)) {
cumuSum2count.put(nextCumuSum, cumuSum2count.get(nextCumuSum) + 1);
}
else {
cumuSum2count.put(nextCumuSum, 1);
}
}
return count;
}
}
提交结果如下:
![](https://i-blog.csdnimg.cn/blog_migrate/14a29a191fc0ca471b370ff33b2a7a48.png)
Python版代码如下:
class Solution:
def subarraySum(self, nums: List[int], k: int) -> int:
numsLen = len(nums)
count = 0
cumulativeSum = [nums[0]]
for i in range(1, numsLen): #计算累加和
cumulativeSum.append(cumulativeSum[-1] + nums[i])
cumuSumDict = {k : 1} #以需要的下一个累加和为key,需要这样的下一个累加和的次数为value
for cumuSum in cumulativeSum:
if cumuSum in cumuSumDict: #当前累加和已在字典中出现,说明当前索引和之前若干个索引之间的区间和为k
count += cumuSumDict[cumuSum]
nextCumuSum = cumuSum + k #想要当前索引和可能的接下来某个索引之间区间和为k,那么其累加和应该在当前基础上+k
if nextCumuSum in cumuSumDict:
cumuSumDict[nextCumuSum] += 1
else:
cumuSumDict[nextCumuSum] = 1
return count
提交结果如下:
![](https://i-blog.csdnimg.cn/blog_migrate/89ab535bc7416d0fa8000b2178e13ebe.png)