1. 程式人生 > >560. Subarray Sum Equals K - Medium

560. Subarray Sum Equals K - Medium

Given an array of integers and an integer k, you need to find the total number of continuous subarrays whose sum equals to k.

Example 1:

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

 

Note:

  1. The length of the array is in range [1, 20,000].
  2. The range of numbers in the array is [-1000, 1000] and the range of the integer k is [-1e7, 1e7].

 

解法:hashmap + prefix sum (cummulative sum)

 

prefix sum:

prefixsum[x] = sum of subarray(0, x) = nums[0] + ... + nums[x] = prefixsum[x-1] + nums[x]

e.g. index      0  1  2

  nums      [1, 1, 1]

  prefix sum [0, 1, 2, 3]    sum of subarray(i, j) = prefixsum[j] - prefixsum[i-1]

用了prefix sum,問題就轉化為 j > i時,求有多少對<i, j>滿足 prefixsum[j] - prefixsum[i-1] = k -> prefixsum[j]  - k = prefixsum[i-1]

 

detail:

用hashmap存 連續子陣列之和(sum = nums[0] + ... + nums[i]) 及其出現次數,以便查詢sum-k出現的次數(有多少prefix sums等於sum - k,nums[0] + ... + nums[j] = sum - k,就有相同數目的prefix sums,nums[j+1] + ... + nums[i] = k)。

遍歷nums中元素,累加當前元素更新sum,在map中查詢key=sum-k出現的次數(即,在遍歷過的陣列元素中,prefixsum[i-1] = prefixsum[j]  - k 出現的次數),然後更新當前prefixsum在map中的次數。

注意:map初始化,map.put(0, 1)!即當sum剛好等於k的時候,從陣列下標為0到當前位置這一子陣列的和就是k,也符合題意。

時間:O(N),空間:O(N)

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