算法思维 ---- 前缀和技巧

前缀和技巧

前缀和技巧 主要用于解决数组的子数组问题。

使用前缀和的经典题目：

• 560.Subarray Sum Equals K
• 930.Binary Subarrays With Sum

1. 什么是前缀和

这里以 LeetCode 的一道题目，560. Subarray Sum Equals K，来讲解前缀和技巧。

题目中给定一个数组 nums 和一个数字 k，找出数组中连续的子数组的和等于给定值 k 的子数组数量。

这道题就是枚举每一组子数组，然后求和，统计等于 k 的子数组数目。但是这样的操作，显然时间复杂度很高，而且可能会超时，在面试中面试官明显不想看到这种代码思路。

这里我们就展示一波 前缀和技巧 吧。

前缀前缀前缀，明显就是当前数组元素 i 的前面的元素的之和的意思嘛。我们开启一个数组长度是 原数组长度+1 的新数组 preSums ，其中每一个元素都是其前面的元素之和，注意不包括当前元素本身。

以 nums = [1, 2, 3] 为例，我们得到的 preSums 就会是 [0, 1, 3, 6]。

你看，前缀和数组就是这么简单，当然做题时也要灵活变通，有时候我们不一定需要开启一个新数组的，只要一个变量来存放操作过程中的前缀和即可。

// 构造前缀和数组
let getPreSum = (nums) => {
  let preSum = [0];
  for (let i = 0; i < nums.length; i++) {
    preSum[i+1] = preSum[i] + nums[i];
  }
  return preSum;
}

既然得到了前缀和数组了，那么这道题只要前缀和之间存在差值为 k 的，就是存在这样的一个子数组了。

let subarraySum = (nums, k) => {
  let perSum = getPreSum(nums);
  
  let count = 0;
  for (let i = 1; i <= nums.length; i++) {
    for (let j = 0; j < i; j++) {
      if (preSum[i] - preSum[j] === k) {
        count++;
      }
    }
  }
  return count;
}

上面的解法的时间复杂度是O(n^2)，空间复杂度是O(n)，n 是nums 的长度。到这里，我们就解决了一道题目了，但是时空复杂度有点高，我们有什么方法优化呢？

2. 优化

上面的解法中，我们有两个for循环，其主要的作用就是找出当前 i 的前缀和，与其前面的前缀和元素 j 之间是否存在差值 k，如果存在，那么元素 j 到 i 的子数组元素之和就是 k ，显然这就是一个结果。

因此，我们可以使用一个 Map 结构，在遍历过程中，如果Map中存在 preSum - k 的键值对，说明存在子数组之和等于 k 的子数组，如果没有则把当前 preSum 存入键值对。

const subarraySum2 = function (nums, k) {
  let map = new Map();
  map.set(0, 1);

  let result = 0;
  let preSum = 0;
  for (let i = 0; i < nums.length; i++) {
    preSum += nums[i];
    result += map.get(preSum - k) ? map.get(preSum - k) : 0;
    map.set(preSum, map.has(preSum) ? map.get(preSum) + 1 : 1);
  }
  return result;
};

注意一下哦，在开启前缀和数组时我们的序号0的元素是不存元素的，默认0，因为存在 (k - 0) 的一个前缀和；所以建立一个 Map 时，我们也需存入 (0,1) 的一个键值对。

至此，我们的时间复杂度降至 O(n) 了。

3. 总结

对于子数组之和的操作，我们都要联想到 前缀和 这个技巧，并且能否使用 HashTable 优化算法。

学习算法，最重要的是思维，是能快速联想到使用什么算法、使用什么结构的反应。