和为 K 的子数组

118 阅读1分钟

560. 和为 K 的子数组 - 力扣(LeetCode)

给你一个整数数组 nums 和一个整数  k ,请你统计并返回 该数组中和为  k  的连续子数组的个数 。

示例 1:

输入: nums = [1,1,1], k = 2
输出: 2

示例 2:

输入: nums = [1,2,3], k = 3
输出: 2

提示:

  • 1 <= nums.length <= 2 * 10^4
  • -1000 <= nums[i] <= 1000
  • -10^7 <= k <= 10^7

思路

本题最简单的方法就是枚举。我们只需要枚举出所有以 i(0<=0 < n)结尾和为 k 的连续子数组个数即可,对于连续子数组[j, ..., i] 进行求和和 k 比较是否相等,如果相等,个数 +1,代码很简单,这里不做过多讲解。

本题还可以用 前缀和 + 哈希表 优化,我们用 preSum 来存储前缀和,preSum[i] = nums[0] + nums[1] + ... + nums[i - 1], 用 map 来存储前缀和 sumi 之前出现的次数。如果 preSum[i] - preSum[j] = k (i > j),则连续子数组[j-1, j, ..., i - 1] 的和为 k,设 start = preSum[i] - kmap[start] 就是 i 之前前缀和为 start 的次数,个数 count += map[start], 并把新的前缀和存入 mapmap[preSum[i]] += 1

解题

/**
 * @param {number[]} nums
 * @param {number} k
 * @return {number}
 */
var subarraySum = function (nums, k) {
  const n = nums.length;
  const preSum = Array(n + 1).fill(0);
  const map = new Map([[0, 1]]);
  let count = 0;
  for (let i = 0; i < n; i++) {
    const sum = preSum[i] + nums[i];
    const diff = sum - k;
    count += map.get(diff) || 0;
    preSum[i + 1] = sum;
    map.set(sum, (map.get(sum) || 0) + 1);
  }
  return count;
};