和为K的子数组

哈希表数组前缀和

Leetcode热题100

发布日期: 2023-05-15

更新日期: 2023-05-15

文章字数: 895

阅读时长: 3 分

描述

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

来源：力扣（LeetCode）
链接：https://leetcode.cn/problems/subarray-sum-equals-k
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

题解

当我们遍历到数组中的第 i 个元素时，我们计算从数组的第一个元素到第 i 个元素的和，即前缀和。假设前缀和为 sum_i。

如果在之前的前缀和中，存在一个前缀和 sum_j，使得 sum_j 等于 sum_i 减去 k，那么说明从数组中第 j + 1 个元素到第 i 个元素的和为 k。这是因为 sum_i 减去 sum_j 等于 k。

我们可以用一个哈希表 preSum 来记录前缀和出现的次数。初始时，前缀和为 0 的出现次数为 1，因为数组中一个元素都不选也可以构成和为 0 的子数组。然后，当我们遍历到数组中的第 i 个元素时，我们先更新前缀和 sum_i，然后判断是否存在一个前缀和 sum_j，使得 sum_j 等于 sum_i 减去 k。如果存在，那么我们就将结果变量 res 加上 preSum[sum_j]，即从数组中第 j + 1 个元素到第 i 个元素的和为 k 的子数组的个数。最后，我们将前缀和 sum_i 的出现次数加 1，以便后续的计算。

整个过程中，我们只需要遍历一遍数组，所以时间复杂度为 O(n)，其中 n 是数组的长度。因为我们用了一个哈希表来记录前缀和出现的次数，所以空间复杂度也是 O(n)。

C++代码

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        // 定义一个哈希表，用于记录前缀和出现的次数
        unordered_map<int, int> preSum;
        // 初始化前缀和为 0 的出现次数为 1
        preSum[0] = 1;
        // 定义变量 sum 和结果变量 res
        int sum = 0, res = 0;
        // 遍历数组中的每个元素
        for(int i = 0; i < nums.size(); i++){
            // 更新前缀和
            sum += nums[i];
            // 如果前缀和中存在 sum - k，说明存在一个子数组的和为 k
            if(preSum.count(sum - k)){
                res += preSum[sum - k];
            }
            // 更新前缀和中 sum 出现的次数
            preSum[sum]++;
        }
        // 返回结果
        return res;
    }
};