Leetcode | 560. 和为 K 的子数组

560. 和为 K 的子数组

文章目录

[560. 和为 K 的子数组](https://leetcode.cn/problems/subarray-sum-equals-k/)
- 题目
- 解法1：暴力枚举
- 解法2：前缀和
- 解法3：[官方题解](https://leetcode.cn/problems/subarray-sum-equals-k/solutions/238572/he-wei-kde-zi-shu-zu-by-leetcode-solution/)

欢迎关注公众号“三戒纪元”

题目

给你一个整数数组 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

解法1：暴力枚举

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        int count = nums.size();
        if(count == 0) {
            return 0;
        }

        int match = 0;
        for(int i = 0; i < count; ++i) {
            int sum = 0;
            for(int j = i; j < count; ++j) {超时
                sum += nums[j];
                if(sum == k) {
                    ++match;
                }
            }
        }
        return match;
    }
};

但是这种解法对于计算包含N个1和N个0的数组会超时

即便是暴力穷举，官方题解思路也别出心裁，虽然同样会超时：

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        int count = 0;
        for (int start = 0; start < nums.size(); ++start) {
            int sum = 0;
            for (int end = start; end >= 0; --end) {
                sum += nums[end];
                if (sum == k) {
                    count++;
                }
            }
        }
        return count;
    }
};

作者：力扣官方题解
链接：https://leetcode.cn/problems/subarray-sum-equals-k/solutions/238572/he-wei-kde-zi-shu-zu-by-leetcode-solution/
来源：力扣（LeetCode）
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。

解法2：前缀和

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        int count = nums.size();
        if(count == 0) {
            return 0;
        }
        sums.resize(count + 1, 0);

        int match = 0;
        for(int i = 0; i < count; ++i) {
            int sum = 0;
            for(int j = i; j >= 0; --j) {
                sum += nums[j];
                if(sum == k) {
                    ++match;
                }
            }
        }
        return match;
    }

    vector<int> sums;
};

同样会超时。

解法3：官方题解

我们可以基于方法一利用数据结构进行进一步的优化，我们知道方法一的瓶颈在于对每个 iii，我们需要枚举所有的 jjj 来判断是否符合条件，这一步是否可以优化呢？答案是可以的。

我们定义 $p re [i]$ 为 $[0.. i]$ 里所有数的和，则 $p re [i]$ 可以由 $p re [i - 1]$ 递推而来，即：

$p re [i] = p re [i - 1] + n u m s [i]$

那么「 $[j .. i]$ 这个子数组和为 k 」这个条件我们可以转化为
$p re [i] - p re [j - 1] == k$
简单移项可得符合条件的下标 jjj 需要满足
$p re [j - 1] == p re [i] - k$
所以我们考虑以 i 结尾的和为 k 的连续子数组个数时只要统计有多少个前缀和为 $p re [i] - k$ 的 $p re [j]$ 即可。我们建立哈希表 $m p$ ，以和为键，出现次数为对应的值，记录 $p re [i]$ 出现的次数，从左往右边更新 $m p$ 边计算答案，那么以 iii 结尾的答案 $m p [p re [i] - k]$ 即可在 $O (1)$ 时间内得到。最后的答案即为所有下标结尾的和为 k的子数组个数之和。

需要注意的是，从左往右边更新边计算的时候已经保证了 $m p [p re [i] - k]$ 里记录的$ pre[j]$的下标范围是 $0 \leq j \leq i$ 。同时，由于 $p re [i]$ 的计算只与前一项的答案有关，因此我们可以不用建立$ pre$ 数组，直接用 $p re$ 变量来记录 $p re [i - 1]$ 的答案即可。

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        int count = nums.size();
        if(count == 0) {
            return 0;
        }
        unordered_map<int, int> mp;
        mp[0] = 1;

        int match = 0;
        int pre = 0;
        for(int &num : nums) {
            pre += num;
            if(mp.find(pre - k) != mp.end() ) {
                match += mp[pre - k];
            }
            mp[pre]++;
        }
        return match;
    }

};