leetcode - [560] Subarray Sum Equals K

2020/06/30

问题描述

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

计算数组中有多少组连续子数组的和为k

举个例子:

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

这道题非常有意思,先来说下我和这道题的故事:

前两天我正在复习HashTable类型的题目,正好复习到 [1] Two Sum ,然后顺带做了类似题目:

[15] 3Sum
[18] 4Sum
[167] Two Sum II - Input array is sorted
[653] Two Sum IV - Input is a BST

然后,还有这道题。

于是,阴差阳错,我通过花花了解到了prefix sum的思想。

花花酱 LeetCode 560. Subarray Sum Equals K - 刷题找工作 EP176

接着,我在网上看了好多prefix sum相关资料:

560 - Subarray Sum Equals K【FLAG高频精选面试题讲解】
Prefix Sum Algorithm | Prefix Sum Array | Difference Array | Range Sum QueryO(1) | EP2
LeetCode 560. Subarray Sum Equals K 中文详解
@shawngao

下面,我就来梳理一下。

解法1

首先,是暴力解法。

来看下示意图:

它的整体思路是这样的:

首先,它借助两个指针ij,不断地变化并计算区间内所有元素的和。

可以从上图看到,长度为3的数组,一共需要计算6次和,并且有很多重复的计算。

来看下实现:

class Solution {
    public int subarraySum(int[] nums, int k) {
        
        int ans = 0;
        int len = nums.length;
        
        for (int i = 0; i < len; i++){
            // Include one element.
            for (int j = i; j < len; j++){
                int sum = 0;
                for (int i = i; i <= j; i++){
                    sum += nums[i];
                }
                if (sum == k){
                    ans++;
                }
            }
        }
        return ans;
    }
}

这种解法超时了,并且时间复杂度为O(n^3)

那么,我们如何改进呢?

再来看下上面的示意图,当i = 0时,我们需要计算三段区间。

你有没有发现,当前区间的和,都依赖于上一个区间的和。

所以,我们可以将每次累加的区间和使用一个变量prefixSum进行保存。

class Solution {
    public int subarraySum(int[] nums, int k) {
        
        int ans = 0;
        int len = nums.length;
        
        for (int i = 0; i < len; i++){
            int prefixSum = 0;
            // Include one element.
            for (int j = i; j < len; j++){
                prefixSum += nums[j];
                if (prefixSum == k){
                    ans++;
                }
            }
        }
        return ans;
    }
}

继续优化。

先来了解一下什么是prefix Sum数组。

来看下示意图:

很简单,所谓prefix Sum数组,就是依次遍历原始数组,累加当前元素与之前的所有元素之和。

注意,我们需要多一个首元素0,便于后续的计算。

那么,我们如何计算nums[0]nums[1]区间内的和?

有没有发现: sum of range[0, 1] = preSum[2] - preSum[0] = 3 - 0 = 3

再来,如何计算nums[0]nums[3]区间内的和?

老规矩: sum of range[0, 3] = preSum[4] - preSum[0] = 5 - 0 = 5

当心咯,如何计算nums[2]nums[3]区间内的和?

来简单推导一下:

sum of range[2, 3]
= sum of range[0, 3] - sum of range[0, 1]
= (preSum[4] - preSum[0]) - (preSum[2] - preSum[0])
= preSum[4] - preSum[2]
= 5 - 3
= 2

好了,有了上面的基础,我们来看下这道题的整体示意图:

来看下实现:

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

Enjoy it !


一位喜欢提问、尝试的程序员

(转载本站文章请注明作者和出处 姚屹晨-yaoyichen

Post Directory