[Leetcode] 560. Subarray Sum Equals K | 和为K的子数组

yPhantom 2020年02月02日 60次浏览

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].

给了一个整数数组nums和一个整数k,让求连续子数组中和为k的子数组个数。

这道题用到了一个叫累加和快速求子集和的方法。即如果我们想快速求出一个数组的子数组的和,可以通过O(n2)的方式,先新建一个累加和数组,然后根据以下代码:

for (int i = 0; i < sums.length; i++) {
  sums[i] // 代表子数组nums[0:i],两边都是闭区间
  for (int j = i - 1; j >= 0; j--) {
    sums[i] - sums[j] // j从i-1开始循环,代表以nums[i]为终点的子数组的和
  }
}

如数组nums为[1, 3, 4, 6],累加和数组sums为[1, 4, 8, 14],假设i = 3,则在累加和方法中能求出[1, 3, 4, 6]、[3, 4, 6]、[4, 6]、[6]这四个子数组之和。

解法一

通过累加和方法,时间复杂度为O(n2)。

public class Solution {
    public int subarraySum(int[] nums, int k) {
        int count = 0;
        int[] sum = new int[nums.length + 1];
        sum[0] = 0;
        for (int i = 1; i <= nums.length; i++)
            sum[i] = sum[i - 1] + nums[i - 1];
        for (int start = 0; start < nums.length; start++) {
            for (int end = start + 1; end <= nums.length; end++) {
                if (sum[end] - sum[start] == k)
                    count++;
            }
        }
        return count;
    }
}

解法二

但是在讨论区中比较推崇是O(n)的HashMap+累加和的方法。

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

我们可以这样理解:

假设j >= i,连续子数组nums[i:j]的和为k,那么nums[0:i - 1]的和即为sums[j] - k。在上面解法中,我们即是以sum - k为hashmap的key值,因为我们是从0开始遍历的,在遍历的过程,如果存在当前i位置之前的累积和,也就是说如果preSum这个map中已经存有sum - k了,假设出现了n次,那么和为k的子数组一定也有n个。

也就是说在当前的i之前,已经有nums[0:p]、nums[0:q]、nums[0:m]的累加和为sum - k了,一共有n个,那么和为k的连续子数组有多少个呢?nums[p + 1, i]、nums[q + 1, i]、nums[m + 1, i],一样有n个。

除了上一点之外,因为sums累加和是从0开始的累加和,而题意要求的是可以是从0开始,也可以是从任意位置开始的累加和。我们初始化一个(0,1) 就是代表sum - k = 0,即从0开始加到当前数正好等于k,默认为1次。

参考资料:

Java Solution, PreSum + HashMap

相似题目:

437. Path Sum III