LeetCode 560.和为K的子数组

题目

给定一个整数数组和一个整数 k，你需要找到该数组中和为 k 的连续的子数组的个数。

示例

输入:nums = [1,1,1], k = 2
输出: 2 , [1,1] 与 [1,1] 为两种不同的情况。

说明 :

• 数组的长度为 [1, 20,000]。
• 数组中元素的范围是 [-1000, 1000] ，且整数 k 的范围是 [-1e7, 1e7]。

题目链接

题目分析

暴力

显然可以用一个时间复杂度为O(n^3)的暴力方法遍历数组，但是时间复杂度过高，会超时。

为了优化这种算法，我们引入前缀和数组。

前缀和

使用前缀和数组可以降低时间复杂度至O(n^2)。所谓前缀和数组，就是ans[i] = A[0] + A[1] + ... + A[i]。

我们只需要遍历前缀和数组即可：

for (auto a : nums) {
    temp += a;
    ans.push_back(a);
}
for (int i = 0; i < ans.size(); i++) {
    temp = 0;
    for (int j = i; j < ans.size(); j++) {
        temp += nums[j];
        if (temp == k) res++;
    }
}
return res;

虽然降低了复杂度至O(n^2)，但还是太慢。由于这道题目只关注次数，而不关注具体是那些子数组的和为K，所以我们可以借助hashmap来进一步降低时间复杂度。

前缀和+hashmap

使用hashmap储存前缀和与该前缀和出现的次数，然后检验前缀和 - k是否在数组内，如果在数组内，则说明可以构成和为k的子数组。

for (auto n : nums){
    sum += n;
    if (m.count(sum - k)) res += m[sum - k];
    m[sum]++;
}

最后返回res即可。

题目解答

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        unordered_map<int, int> m;
        m[0] = 1;
        int sum = 0;
        int res = 0;
        for (auto n : nums){
            sum += n;
            if (m.count(sum - k)) res += m[sum - k];
            m[sum]++;
        }
        return res;
    }
};