Leetcode | 560. 和为 K 的子数组

560. 和为 K 的子数组

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

题目

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

示例 1：

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

示例 2：

输入：nums = [1,2,3], k = 3
输出：2
1
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;
    }
};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

但是这种解法对于计算包含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）
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

解法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;
};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

同样会超时。

解法3：官方题解

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

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

$$pre[i]=pre[i-1]+nums[i]$$

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

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

};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23