560. 和为 K 的子数组 - 前缀和思想
560. 和为 K 的子数组 - 前缀和思想
在算法题中,前缀和是一种能快速计算 “数组中某段连续元素之和” 的预处理方法,核心思路是 “提前计算并存储中间结果,避免重复计算”
前缀和的定义:
对于一个数组 nums,我们可以创建一个 “前缀和数组” prefix,其中 prefix[i] 表示 “数组 nums 中前 i 个元素的和”(注意:这里的 “前 i 个” 通常从第 1 个元素开始算)。
以上面的步数数组为例:
- prefix[0] = 0(约定的初始值,方便计算)
- prefix[1] = 第1天步数 = 3
- prefix[2] = 第1天 + 第2天 = 3 + 1 = 4
- prefix[3] = 第1天 + 第2天 + 第3天 = 3 + 1 + 4 = 8
- prefix[4] = 前4天总和 = 3 + 1 + 4 + 2 = 10
- prefix[5] = 前5天总和 = 3 + 1 + 4 + 2 + 5 = 15
所以前缀和数组是 [0, 3, 4, 8, 10, 15]。
前缀和的妙用:快速算 “区间和”有了前缀和数组,计算 “从第 l 天到第 r 天的总步数” 时,直接用:
区间和 = prefix [r] - prefix [l-1]
比如:
- 第 2 天到第 4 天:l=2,r=4 → prefix[4] - prefix[1] = 10 - 3 = 7(和之前手动算的一致)
- 第 1 天到第 3 天:l=1,r=3 → prefix[3] - prefix[0] = 8 - 0 = 8(正确)
题目描述
https://leetcode.cn/problems/subarray-sum-equals-k/description/?envType=study-plan-v2&envId=top-100-liked
给你一个整数数组 nums 和一个整数 k ,请你统计并返回 该数组中和为 k 的子数组的个数 。
子数组是数组中元素的连续非空序列。示例 1:
输入:nums = [1,1,1], k = 2输出:2
示例 2:
输入:nums = [1,2,3], k = 3输出:2提示:
1 <= nums.length <= 2 * 104
-1000 <= nums[i] <= 1000
-107 <= k <= 107
题解
class Solution {public int subarraySum(int[] nums, int k) {int len = nums.length;int preSum = 0; // 存储前缀和HashMap<Integer,Integer> map = new HashMap(); // 用来存储前缀和出现的次数map.put(0,1); // 初始化// 若某一个前缀和preSum恰好等于k,此时preSum - k = 0,能直接匹配到初始值 1。int res = 0;for(int x: nums){preSum += x; // 计算前缀和/*假设当前前缀和是preSum(对应数组前 i 个元素的和)我们要找 “和为 k 的子数组”,也就是找两个前缀和 preSum[i] 和 preSum[j](j < i),使得:preSum[i] - preSum[j] = k变形后就是:preSum[j] = preSum[i] - k需要找到之前preSum[j]出现过几次,那么就有几个子数组和为k*/res += map.getOrDefault(preSum - k, 0);// 更新前缀和出现的次数map.put(preSum, map.getOrDefault(preSum, 0) + 1);}return res;}
}
核心问题:如何用前缀和表示 “子数组的和”?
任意一个子数组的和,都可以表示为 “两个前缀和的差”。
比如:
- 子数组 [1](索引 1)的和 = 1对应:preSum[2] - preSum[1] = 3 - 2 = 1
- 子数组 [3](索引 2)的和 = 3对应:preSum[3] - preSum[2] = 6 - 3 = 3
- 子数组 [2,1](索引 0-1)的和 = 3对应:preSum[2] - preSum[0] = 3 - 0 = 3
关键逻辑:“子数组和为 k” 等价于什么?
我们要找 “和为 k 的子数组”,也就是找两个前缀和 preSum[i] 和 preSum[j](j < i),使得:
preSum[i] - preSum[j] = k
变形后就是:
preSum[j] = preSum[i] - k
这句话的意思是:如果当前的前缀和是 preSum[i],那么只要之前出现过等于 preSum[i] - k 的前缀和(记为 preSum[j]),那么子数组 nums[j…i-1] 的和就一定是 k。