使用「前缀和」解决力扣 560 号问题：和为 K 的子数组

前缀和，是指一个数组的某下标 i 之前的所有数组元素的和（包不包括 i 都可）。比如，现有数组：[1, 2, 3, 4, 5]，那么对应的前缀和数组就是 [0, 1, 3, 6, 10, 15]。使用前缀和可以通过预处理数组来降低每次求和的复杂度。

如果没有前缀和数组，我们想计算数组位置 [1, 3] 的和，是不是需要写一个 for(int i = 1; i <= 3; i++) sum += nums[i]; 循环。计算一次还好，如果是需要计算上千上万次呢？如果数组非常大呢？这样简单的用 for 循环是不是很费时间？

我们通过预处理计算出前缀和数组，就能很容易的在 O(1) 的时间复杂度内求得任意数组位置的和。

int[] nums = new int[]{1,2,3,4,5};
int[] prefixSum = new int[nums.length + 1];

prefixSum[0] = 0;

for (int i = 0; i < nums.length; i++) {
    prefixSum[i + 1] = nums[i] + prefixSum[i];
}

现在我们计算任意数组位置 [l, r] 的和：prefixSum[r + 1] - prefixSum[l]。假如计算 [1, 3] 的和就是：prefixSum[3 + 1] - prefixSum[1]。

现在看一下力扣 560 号问题：和为 K 的子数组。题目让你计算和为 k 的连续子数组的和。

很容易想到使用前缀和的暴力解法：

int ans = 0;
int n = nums.length;
int[] prefixSum = new int[n + 1];

prefixSum[0] = 0;

for (int i = 0; i < n; i++) {
    prefixSum[i + 1] = nums[i] + prefixSum[i];
}
    
for (int i = 1; i <= n; i++) {
    for (int j = 0; j < i; j++) {
        if (prefixSum[i] - prefixSum[j] == k) ans++;
    }
}

return ans;

接下来该怎么优化呢？

从内层循环下手。先来看下内层循环的作用：循环 i 次，找出符合 if 条件的组合一共有多少个。显然对于 if 条件，只有 j 是一直在变的，所以换句话说就是：找出符合 if 条件的 j 一共有多少个。

将 if 条件移项得：if (prefixSum[j] == prefixSum[i] - k)。那么现在只要知道有几个 prefixSum[j] 和 prefixSum[i] - k 相等，就可以避免内层 for 循环。我们还可以使用一个 HashMap 来记录和 prefixSum[i] - k 相等的 prefixSum[j] 有几个，这样内存循环就可以直接省掉了。

另外，prefixSum[i] 是不是就是 nums[0...i]的和？因此可以使用一个变量 sum0_i 来承接。这样连前缀和数组也省掉了，就只剩前缀和思想在这了。

完整代码：

int ans = 0;
HashMap<Integer, Integer> sumMap = new HashMap<>();
int sum0_i = 0;
    
sumMap.put(0, 1);
    
for (int i = 0; i < nums.length; i++) {
    sum0_i += nums[i];
    
    if (sumMap.containsKey(sum0_i - k)) {
        ans += sumMap.get(sum0_i - k);
    }
    
    sumMap.put(sum0_i, sumMap.getOrDefault(sum0_i, 0) + 1);
}
    
return ans;