[LeetCode] 862. Shortest Subarray with Sum at Least K 和至少为K的最短子数组

Return the length of the shortest, non-empty, contiguous subarray of `A` with sum at least `K`.

If there is no non-empty subarray with sum at least K, return -1.

Example 1:

Input: A = [1], K = 1
Output: 1

Example 2:

Input: A = [1,2], K = 4
Output: -1

Example 3:

Input: A = [2,-1,2], K = 3
Output: 3

Note:

1. 1 <= A.length <= 50000
2. -10 ^ 5 <= A[i] <= 10 ^ 5
3. 1 <= K <= 10 ^ 9

这道题给了我们一个非空整数数组和一个正整数K，让找出非空的子数组使得其和至少为K，找不到的话返回 -1。这道题的难点在于数组中可能有负数，这样的话子数组之和就不会随着长度的增加而增加，从而贪婪算法可能会失效。当然，直接暴力搜索搜索所有的子数组是对 Hard 题目的不尊重，会被 OJ 教育。对于子数组之和的问题，十有八九是要建立累加和数组的，因为其可以快速的计算任意区间和，但即便是有了累加和数组，遍历所有区间和还是会超时。用累加数组计算任意区间 [i, j] 的累加和是用 [0, j] 区间和减去 [0, i-1] 区间和得到的，只有两个区间和差值大于等于K的时候，才会更新结果，所有小于K的区间差是不需要计算的。这样的话，假如能使得所有区间和按照从小到大的顺序排列，那么当前区间和按顺序减去队列中的区间和，一旦差值小于K了，后面的区间和就不用再检验了，这样就可以节省很多运算。

思路有了，下面就来解题吧。这里用一个最小堆，里面放一个数对儿，由区间和跟其结束位置组成。遍历数组中所有的数字，累加到 sum，表示区间 [0, i] 内数字和，判断一下若 sum 大于等于K，则用 i+1 更新结果 res。然后用一个 while 循环，看 sum 和堆顶元素的差值，若大于等于K，移除堆顶元素并更新结果 res。循环退出后将当前 sum 和i组成数对儿加入最小堆，最后看若结果 res 还是整型最大值，返回 -1，否则返回结果 res，参见代码如下：

解法一：

class Solution {
public:
    int shortestSubarray(vector<int>& A, int K) {
        int n = A.size(), res = INT_MAX, sum = 0;
        priority_queue<pair<int, int>, vector<pair<int, int>>, greater<pair<int, int>>> pq;
        for (int i = 0; i < n; ++i) {
            sum += A[i];
            if (sum >= K) res = min(res, i + 1);
            while (!pq.empty() && sum - pq.top().first >= K) {
                res = min(res, i - pq.top().second);
                pq.pop();
            }
            pq.push({sum, i});
        }
        return res == INT_MAX ? -1 : res;
    }
};

我们也可以不用最小堆，直接利用 TreeMap 的自动排序功能，建立区间和跟其结束位置之间的映射。首先建立累加和数组，然后遍历累加和数组中的每一个数字，在 TreeMap 中二分查找第一个大于 sums[i]-K 的 位置 pos，这样所有前面较小的位置x，都有 sum[i]-x >= K 成立，这样只要从开头遍历到 pos 位置，将每个区间长度更新结果 res 即可。更新完成后，要将开头到位置 pos 之间的映射全部删掉，因为这些已经计算过了，就像上面解法中将堆顶元素移除的操作一样，然后建立当前区间和跟结束位置i之间的映射，参见代码如下：


解法二：

class Solution {
public:
    int shortestSubarray(vector<int>& A, int K) {
        int n = A.size(), res = INT_MAX;
        map<int, int> sumMap;
        vector<int> sums(n + 1);
        for (int i = 1; i <= n; ++i) sums[i] = sums[i - 1] + A[i - 1];
        for (int i = 0; i <= n; ++i) {
            auto pos = sumMap.upper_bound(sums[i] - K);
            for (auto it = sumMap.begin(); it != pos; ++it) {
                res = min(res, i - it->second);
            }
            sumMap.erase(sumMap.begin(), pos);
            sumMap[sums[i]] = i;
        }
        return res == INT_MAX ? -1 : res;
    }
};

上面两种解法虽然都可以通过 OJ，但也仅仅是险过，来看一种时间和空间的击败率都很高的方法，这里用到了双向队列 deque，这是一种两头都能操作的飞起的数据结构。双向队列不像优先队列那样自动排序，这样就节省了排序的时间，我们是按照数组原顺序将数字下标加入双向队列的。在建立好累加和数组之和，遍历其每个累加和，然后用一个 while 循环，从双向队列的开头开始遍历，假如区间和之差大于等于K，就移除队首元素并更新结果 res。之后这个 while 循环非常重要，能有这么高的击败率，全要靠这个循环，这个是从双向队列的末尾开始往前遍历，假如当前区间和 sums[i] 小于等于队列末尾的区间和，则移除队列末尾元素。这是为啥呢？因为若数组都是正数，那么长度越长，区间和一定越大，则 sums[i] 一定大于所有双向队列中的区间和，但由于可能存在负数，从而使得长度变长，区间总和反而减少了，之前的区间和之差都没有大于等于K，现在的更不可能大于等于K，这个结束位置可以直接淘汰，不用进行计算。循环结束后将当前位置加入双向数组即可，参见代码如下：


解法三：

class Solution {
public:
    int shortestSubarray(vector<int>& A, int K) {
        int n = A.size(), res = INT_MAX;
        deque<int> dq;
        vector<int> sums(n + 1);
        for (int i = 1; i <= n; ++i) sums[i] = sums[i - 1] + A[i - 1];
        for (int i = 0; i <= n; ++i) {
            while (!dq.empty() && sums[i] - sums[dq.front()] >= K) {
                res = min(res, i - dq.front());
                dq.pop_front();
            }
            while (!dq.empty() && sums[i] <= sums[dq.back()]) {
                dq.pop_back();
            }
            dq.push_back(i);
        }
        return res == INT_MAX ? -1 : res;
    }
};

Github 同步地址:

https://github.com/grandyang/leetcode/issues/862

参考资料：

https://leetcode.com/problems/shortest-subarray-with-sum-at-least-k/

https://leetcode.com/problems/shortest-subarray-with-sum-at-least-k/discuss/143726/C%2B%2BJavaPython-O(N)-Using-Deque

https://leetcode.com/problems/shortest-subarray-with-sum-at-least-k/discuss/143768/C%2B%2B-solution-by-using-priority_queue

https://leetcode.com/problems/shortest-subarray-with-sum-at-least-k/discuss/153515/c%2B%2B-upper_bound-%2B-prefix-sum-easy-to-remember-and-impl-in-10mins

[LeetCode All in One 题目讲解汇总(持续更新中...)](https://www.cnblogs.com/grandyang/p/4606334.html)