前缀和算法疑问：为何初始化obj为{0:1}？obj[sum-k]作用是什么？

问题描述 #

给定整数数组nums和整数k，返回总和等于k的子数组（数组内连续非空元素序列）的总数。

实现代码 #

var subarraySum = function(nums, k) {
  let obj = {
    0: 1
  };
  let count = 0;
  let sum = 0;
  for (let i = 0; i < nums.length; i++) {
    sum += nums[i];
    if (obj[sum - k]) {
      count += obj[sum - k];
    }
    obj[sum] = ++obj[sum] || 1;
  }

  return count;
};

console.log(subarraySum([1, 2, 3, 3, 2, 1, 3], 4)) // 输出：2
console.log(subarraySum([1, 2, 3, 3, 2, 1, 3], 5)) // 输出：3

疑问解答 #

1. 为何要将obj初始化为{0:1}？ #

这是为了覆盖从数组第一个元素开始的子数组和恰好等于k的情况。

用前缀和思路解释：遍历到第i个元素时，当前前缀和sum如果刚好等于k，那么sum - k = 0。如果obj没有0:1这个初始记录，就会漏掉这种直接符合条件的子数组。

举个实际例子：数组[1,3]，k=4。遍历到第二个元素时，sum=4，sum -k=0，此时obj里的0:1会让我们把[1,3]这个符合条件的子数组计入总数。

2. obj[sum - k]这一操作具体实现了什么逻辑？ #

obj是用来记录前缀和出现次数的哈希表，sum是当前遍历到第i个元素的前缀和（即从数组开头到第i个元素的累加和）。

我们的核心目标是找：有没有之前出现过的前缀和prevSum，满足sum - prevSum = k——也就是prevSum = sum -k。如果存在这样的prevSum，说明从prevSum对应的位置的下一个元素到当前i的元素，这一段连续子数组的和正好是k。

obj[sum -k]的值就是prevSum出现的次数，每出现一次就对应一个符合条件的子数组，所以把这个次数加到count里，就能累计所有符合要求的子数组数量。

比如数组[3,1,2]，k=3：

• 遍历第一个元素时，sum=3，sum -k=0，obj里有0:1，count加1（对应子数组[3]）；
• 遍历第三个元素时，sum=6，sum -k=3，obj里3出现过1次，count再加1（对应子数组[1,2]）；
• 最终count=2，结果正确。

内容的提问来源于stack exchange，提问作者Aashrun Gautam