leetcode-和为K的子数组

题目来自LeetCode，链接：subarray-sum-equals-k。具体描述为：给定一个整数数组和一个整数 k，你需要找到该数组中和为 k 的连续的子数组的个数。

示例:

输入:nums = [1,1,1], k = 2
输出: 2 , [1,1] 与 [1,1] 为两种不同的情况。

一个很容易想到的思路就是暴力法，直接遍历所有可能的子数组，最好也只能做到时间复杂度为$O(N^2)$，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)$。

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;
    }
}

提交结果如下:

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

提交结果如下: