#560 Subarray Sum Equals K

Description

Given an array of integers nums and an integer k, return the total number of continuous subarrays whose sum equals to k.

Examples

Example 1:

Input: nums = [1,1,1], k = 2
Output: 2

Example 2:

Input: nums = [1,2,3], k = 3
Output: 2

Constraints:

1 <= nums.length <= 2 * $10^4$
-1000 <= nums[i] <= 1000
-$10^7$ <= k <= $10^7$

思路

最开始的时候我想用的是类似于KMP的算法，类似于动态窗口的方式，一点点往前挪动

class Solution {
    public int subarraySum(int[] nums, int k) {
        int answer = 0;
        int pointer = 0;
        int tempSum = 0;
        
        for(int i = 0; i < nums.length; i++){
            tempSum += nums[i];
            
            while (tempSum > k){
                tempSum -=nums[pointer];
                pointer ++;
            }
            
            if (tempSum == k && i - pointer >= 0){
                answer ++;
                tempSum -=nums[pointer];
                pointer ++;
            }
        }
        
        return answer;
    }
}

但貌似这个方法只适用于全正数的情况，对下面这种样例就wrong answer了

[-1, -1, 1]
0

然后就想不出了，老老实实用了暴力求解 O(n^2)

class Solution {
    public int subarraySum(int[] nums, int k) {
        int answer = 0;
        int tempSum = 0;
        
        for(int i = 0; i < nums.length; i++){
            tempSum = 0;
            for(int j = i; j < nums.length; j++){
                tempSum += nums[j];
                if(tempSum == k){
                    answer ++;
                }
            }
        }
        return answer;
    }
}

居然能ac！我还以为会TLE，虽然时间上是1227ms

但是这道题肯定有tricky的解法！翻了下discussion，就是用Map（比较类似于 #1 two sum）。这道题的map用的映射关系是

• $[0,i]$ 的和 (记为 preSum) → 和为 preSum 的 $i$ 的个数

这样有什么好处呢？
假设我们现在位于位置 $j$，且 $[0,j]$ 和为 sum。我们的目标是寻找一段连续的子串，让这个连续子串的和为 $k$，且已知 $[0,j]$ 和为 sum。那么到节点 $j$ 为止，满足连续子串和为 $k$ 的个数就是

• map [sum - k]

不难理解，存在多少个 $i$，就存在多少个满足条件的连续子串。

第一个问题解决了，接下来是可能碰到的第二个问题，怎么处理第一个碰到的满足条件的连续子串？
这种情况下 sum - k = 0，但是前面的计数过程中 [可能] 并没有出现 map[0] 的情况，因此我们在map中需要初始化一个原始情况

• map [0] = 1

也就是将开始处理前的状态所拥有的值置为1

理论讲到这里，可以举个例子。

[-1, -1, 1, 0]
0

• 初始化 map[0] = 1，sum = 0，nums = [-1, -1, 1, 0]，count = 0
• 处理第 1 个数 -1
  • sum[0, 1] = 1
  • sum - k = -1
  • map[sum-k] 没有值，意味着没有 $[0,i]$ 的子串符合条件，所以为0
    • count += 0（count = 0）
  • map[-1] = map[-1] + 1 = 0 + 1 = 1
• 处理第 2 个数 -1
  • sum[0, 2] = -2
  • sum - k = -2
  • map[sum-k] 没有值，意味着没有 $[0,i]$ 的子串符合条件，所以为0
    • count += 0（count = 0）
  • map[-2] = map[-2] + 1 = 0 + 1 = 1
• 处理第 3 个数 1
  • sum[0, 3] = -1
  • sum - k = -1
  • map[sum-k] = 1，有值了，说明前面有 1 个 $[0,i]$ 子串（$i=1$）符合条件
    • count += map[sum-k] （count = 1）
  • map[-1] = map[-1] + 1 = 1 + 1 = 2
• 处理第 4 个数 0
  • sum[0, 4] = -1
  • sum - k = -1
  • map[sum-k] = 2，有值了，说明前面有 2 个 $[0,i]$ 子串（$i=1$， $i=3$）符合条件
    • count += map[sum-k] （count = 3）
  • map[-1] = map[-1] + 1 = 2 + 1 = 3

最后答案就是 count = 3

代码

class Solution {
    public int subarraySum(int[] nums, int k) {
        Map<Integer, Integer> map = new HashMap<>();
        map.put(0, 1);
        
        int sum = 0, count = 0;
        
        for(int i = 0; i < nums.length; i++) {
            sum += nums[i];
            
            if(map.containsKey(sum - k)) {
                count += map.get(sum-k);
            }
            
            map.put(sum, map.getOrDefault(sum, 0) + 1);
        }

        return count;
    }
}