实用算法解析 - 前缀和

简介

前缀和的思路在力扣很多题目的都出现过，常用于处理连续子数组类型的的问题。接下来将用逐层深入的方式来进行介绍，

从一道例题说起

看一道例题(leetcode 560)。

给定一个整数数组和一个整数 k，你需要找到该数组中和为 k 的连续的子数组的个数。

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

说明 :
数组的长度为 [1, 20,000]。
数组中元素的范围是 [-1000, 1000] ，且整数 k 的范围是 [-1e7, 1e7]。

基本思路

首先看到这道题，最容易想到的是暴力破解：求出所有连续子数组的和，然后遍历它们并且统计其中和为k的项数。

为了方便 我们定义一个sum函数 用于求任意连续子数组的和，代码如下：

// 思路1：求出所有连续子数组的和 并统计满足和为k的项数
var subarraySum = function(nums, k) {
    const len = nums.length;
    let count = 0;
    for(let i = 0; i < len; i++){
        for(let j = i+1; j < len; j++){
            if(sum(nums,i,j)===k){
                count++
            }
        }
    }
    return count;
};

// 求数组中从下标startIndex到下标endIndex之间所有元素的和
function sum(arr, startIndex, endIndex){
    let res =0;
    for(let i = startIndex; i<=endIndex ;i++){
        res += arr[i]; 
    }
    return res;
}

这种解法的时间复杂度显然太高了：最外层外面有两层循环 复杂度为O(n^2), sum(arr, i, j)的复杂度为n，所以总的时间复杂度达到了O(n^3)。因此我们要考虑其他的思路。

优化1：引入前缀和

首先我们先想办法优化掉sum函数，因为这个函数对任何一个连续子数组都重新计算一次元素之和，不能有效利用之前的结果。

因此我们引入一个preSum数组，其中preSum[i]表示从数组从开始到下标为i的所有元素的和.也就是:

var getPreSum = function(nums){
    let count = 0;
    const preSum = [];// preSum【i】 表示从开始到第i个元素之和 
    for(let i = 0; i < nums.length; i++){
        if(i === 0){
            preSum[i] = nums[i];
        } else{
            preSum[i] = preSum[i-1] + nums[i]
        }
    }
    preSum[-1] = 0; // 由于数组从第0项开始，所以preSum[-1]表示没有元素 自然是为0. 这是为了方便后面的求解
    return preSum;
}
getPreSum([1,1,1]); // 现在对于题目中输入的[1, 1, 1]数组 我们就得到了一个值为[1,2,3] 的preSum数组

会发现，得到preSum这个数组之后，求解第i项到第j项的元素之和， 等价于求preSum[j] - preSum[i-1]。 （i=0 时，i-1= -1，这也是上面设置preSum[-1] =0的原因）

所以我们现在成功把sum函数去掉了，由于preSum[j+1]- preSum[i]的复杂度是O(1)，所以整体的复杂度也从O(n^3)降低到O(n^2)。使用前缀和之后，我们的代码变成现在这样：

var subarraySum = function(nums, k) {
    const len = nums.length;
    let count = 0;
    const preSum = getPreSum(nums);
    // 请注意这里的i jd的范围和边界条件， 当i = j时, preSum[j] - preSum[i-1] = nums[i]
    for(let i = 0; i < len; i++){
        for(let j = i; j < len; j++){
            if(preSum[j] - preSum[i-1] === k){
                count++；
            }
        }
    }
    return count;
};

优化2：去掉非必须的嵌套循环

但是这样的复杂度O(n^2)依然是不够的，所以我们现在继续优化，目前主要的复杂度集中在嵌套的for循环里，所以先观察下这个循环：
内层循环的关键条件语句是preSum[j] - preSum[i-1] === k,根据等式的基本原理，移项可得 preSum[i-1] === preSum[j] - k，

现在关键点来了，从j的角度考虑(要考虑任意的i<=j 这就是前面提醒读者注意边界条件的原因)：

• 当j = 0时，我们要比较: preSum[-1] 是否等于 preSum[0] - k;
• 当j = 1时，我们要比较: preSum[0] 是否等于 preSum[1] - k, preSum[-1] 是否等于 preSum[1] - k;
• 当j = 2时，我们要比较: preSum[1] 是否等于 preSum[2] - k,preSum[0] 是否等于 preSum[2] - k, preSum[-1] 是否等于 preSum[2] - k;

...

发现了吗？在上述过程中，其实我们多次用到了preSum[0], preSum[1], ... preSum[len] 所以如果我们直接把preSum[i](0<=i<=len-1)缓存起来，就可以解开内层循环了. 所以我们可以考虑用一个hash结构(在js里面通常用obj或者es6里的map)来保存preSum[i] - k， key => value 分别对应 preSum[i] - k => 出现次数。

由于我们预设了preSum[-1] = 0，所以hash结构的第一项默认就是, 0=>1 表示前缀和为0的情况已经出现了一次

接下来，遍历数组中的每一项,并且执行：

1. 查看： 查看现有的hash，是否存在满足[当前的前缀和] - k =[已有的前缀和]
2. 更新： 更新hash,把当前的前缀和添加到hash中去 -- 如果已经存在hash[当前前缀和]，那么出现次数加1; 如果还不存在，那出现次数设置为1；

所以我们可以把上面的思路用代码表示出来：

var subarraySum = function(nums, k) {
    const len = nums.length;
    let count = 0;
    const hash = new Map();
    hash.set(0,1); //预设了preSum[-1]= 0;
    const preSum = getPreSum(nums);
    for(let i = 0; i < len; i++){
        // 操作1： 判断之前出现的前缀和中 是否已经有满足【当前前缀和】=【之前前缀和】- k的项
        const key = preSum[i] - k;        
        if(hash.has(key)){
            count +=  hash.get(key);
        }

        // 操作2：把当前项对应的前缀和放入hash， 这个和上面的操作1的执行顺序是不可以相反的，否则会出现重复计数的问题 可以思考下为什么
        if(!hash.has(preSum[i])){
            hash.set(preSum[i], 1);
        } else {
            hash.set(preSum[i], hash.get(preSum[i]) + 1);
        }

    }
    return count;
};

优化3：取出非必要的preSum结构

经过第二步骤的优化之后，其实已经得到了一个比较好的前缀和算法，只有一层循环，所以时间复杂度为O(n)，需要一个preSum的数组空间和一个map，空间复杂度为O(n)。不过还是有地方可以优化：
preSum必须存在吗？
答案是没有必要的，我们发现getPreSum的本质实质上也是对nums做了一次单层循环，并且在subarraySum函数里， 遍历到i时，我们只需要当前对应的preSum[i]即可**
所以可以改写成以下形式：

var subarraySum = function(nums, k) {
    const len = nums.length;
    let count = 0;
    const hash = new Map();
    hash.set(0,1); //预设了preSum[-1]= 0;
    // const preSum = getPreSum(nums); //这一行不再需要了 用一个临时变量代替
    let currentSum = 0; // 这个初始值其实对应的就是原先的preSum[-1]
    for(let i = 0; i < len; i++){
        currentSum += num[i]； //这一步就求解了preSum[i]

        // 操作1： 判断之前出现的前缀和中 是否已经有满足【当前前缀和】=【之前前缀和】- k的项
        const key = currentSum - k;        
        if(hash.has(key){
            count +=  hash.get(key);
        }

        // 操作2：把当前项对应的前缀和放入hash， 这个和上面的操作1的执行顺序是不可以相反的，否则会出现重复计数的问题 可以思考下为什么
        if(!hash.has(currentSum)){
            hash.set(currentSum, 1);
        } else {
            hash.set(currentSum, hash.get(currentSum) + 1);
        }

    }
    return count;
};

到这里基本上完整的算法就介绍完了。

相关题型

974. 和可被 K 整除的子数组

示例：

输入：A = [4,5,0,-2,-3,1], K = 5
输出：7
解释：
有 7 个子数组满足其元素之和可被 K = 5 整除：
[4, 5, 0, -2, -3, 1], [5], [5, 0], [5, 0, -2, -3], [0], [0, -2, -3], [-2, -3]

学完本文之后有兴趣的可以在leetcode上拿这道类似的题目练练手。

小结

本文针对前缀和算法，以leetcode的一道题为例，按照由浅入深的方式,层层递进地进行介绍。

思路1是我们接触算法较少时最常见最直接的解法；优化1是引入前缀和概念，优化2是保证前缀和时间复杂度满足一般要求的关键步骤，初次理解有难度；优化3则是额外的细节，没有一开始就直接介绍最终的算法，保留中间的轨迹是为了能更方便大家理解。

惯例：如果内容有错误的地方欢迎指出（觉得看着不理解不舒服想吐槽也完全没问题）；如果有帮助，欢迎点赞和收藏，转载请征得同意后著明出处，如果有问题也欢迎私信交流，主页有邮箱地址

---

最后顺便打个小广告：

• RingCentral厦门地区目前有大量hc
• 纯美资外企,有工作优生活（5点多下班 个人时间超长 可以随心所欲撸猫撸厨房撸算法）
• 海景办公，零食水果，节日福利多多
• 免费英文口语课，硅谷工作机会，入职享受10天起超长带薪年假
• 需要内推请私信或投递简历到邮箱[e-mail]
• 全程跟进面试进度，提供力所能及的咨询帮助~