Leetcode（560）——和为 K 的子数组

Leetcode（560）——和为 K 的子数组

题目

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

示例 1：

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

示例 2：

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

提示：

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

题解

关键：与 Leetcode-304 不同，这次只有一次获取，所以像之前一样计算好全部结果的方法不适用

方法一：暴力枚举

思路

​​  考虑以 $i$ 结尾和为 $k$ 的连续子数组个数，则需要统计符合条件的下标 $j$ 的个数，其中 $0\le j\le i$ 且 $[j..i]$ 这个子数组的和恰好为 $k$。

​​  我们可以枚举 $[0..i]$ 里所有的下标 $j$ 来判断是否符合条件，可能有读者会认为假定我们确定了子数组的开头和结尾，还需要 $O(n)$ 的时间复杂度遍历子数组来求和，那样复杂度就将达到 $O(n^3)$ 从而无法通过所有测试用例。但是如果我们知道 $[j,i]$ 子数组的和，就能 $O(1)$ 推出 $[j-1,i]$ 的和，因此这部分的遍历求和是不需要的，我们在枚举下标 $j$ 的时候已经能 $O(1)$ 求出 $[j,i]$ 的子数组之和。

代码实现

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;
    }
};

复杂度分析

时间复杂度：$O(n^2)$，其中 $n$ 为数组的长度。枚举子数组开头和结尾需要 $O(n^2)$ 的时间，其中求和需要 $O(1)$ 的时间复杂度，因此总时间复杂度为 $O(n^2)$。
空间复杂度：$O(1)$

方法二：哈希表 + 前缀和

思路

​​  在这里我们再次引入前缀和的定义：前缀和指 $nums$ 的第 0 项到当前项的和。

​​  但是因为这次只有一次获取，所以与 Leetcode-304 不同，像之前一样计算好全部结果的方法不适用，会导致很大的时间复杂度。所以我们可以像之前的 Leetcode-303 一样找到它们之中的规律或公式来解答问题。
​​  经过观察，可以发现如果有前缀和 $Presum(0,a)$ 的值等于另一个前缀和 $Presum(0,b)$ 减去要找到的值 $k$，且其中数组的下标 $a<b$，则子数组 $sum(a+1,b)=k$。
​​  公式如下：
$$k=Presum[j]-Presum[i-1]=nums[i]+\dots +nums[j]，其中j>i$$

​​  问题：我们发现先计算完前缀和，再一边从头遍历前缀和数组，一边根据公式查询需要的前缀和是否存在，同时还需要判断该前缀和的范围（即 (0,n) 的 n ）是否大于当前遍历到的前缀和（即 n 是否大于 i）。这就导致了很高的时间复杂度，与我们使用前缀和的初衷（减少大量的时间复杂度）相矛盾。
​​  解决方法：在构建前缀和时，就查找是否存在与 包含当前项的前缀和 满足公式的 前缀和（单独一个包含当前项的前缀和也是其子数组）——这就避免了判断找到的前缀和是否是在当前项的之前的前缀和（因为后面的前缀和都还没计算出来）。
​​  同时为了能快速查询到当前项的前缀和所需要的值对应的子前缀和有几个，我们选择使用哈希表。unordered_map<前缀和, 个数>

代码实现

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        int ans = 0, presum = 0;    // ans是最后要返回的结果，presum是包含当前项的前缀和
        unordered_map<int, int> Presums;    // unordered_map<前缀和, 个数> 表示不包含当前项的前缀和以及对应个数
        for(auto& it: nums){
            presum += it;           // 更新包含当前项的前缀和
            if(!Presums.empty())	// 第一次时不存在不包含当前项的前缀和
                if(Presums.count(presum - k) != 0)
                    ans += Presums[presum - k];
            if(presum == k)
                ans++;
            if(Presums.count(presum) == 0)     // 更新 Presums 
                Presums.emplace(presum, 1);
            else Presums[presum]++;
        }
        return ans;
    }
};

复杂度分析

时间复杂度：我们遍历数组的时间复杂度为 $O(n)$，中间利用哈希表查询删除的复杂度均为 $O(1)$，因此总时间复杂度为 $O(n)$，其中 $N$ 是数组的元素个数。
空间复杂度：$O(n)$，其中 $n$ 为数组的长度。哈希表在最坏情况下可能有 $n$ 个不同的前缀和作为键值，因此需要 $O(n)$ 的空间复杂度。