【Leetcode】862. Shortest Subarray with Sum at Least K

题目地址：

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

给定一个数组$A$，求其最短的和大于等于$K$的子数组的长度。题目保证$K>0$。

法1：单调栈 + 二分。首先要求子数组的和，想到前缀和。设数组$s$是$A$的前缀和，具体来说，规定$s[0]=0$，$s[i]=A[0]+...+A[i-1]$。那么问题转为求在$s[y]-s[x]\ge K$的情况下的最小的$y-x$。首先，在固定子数组右端点$i$的情况下，考虑左端点，如果$x<y$并且$s[x]\ge s[y]$的话，那$x$一定不可能是左端点，因为$s[i]-s[y]\ge s[i]-s[x]$，左端点选$x$不如选$y$。所以在固定右端点的情况下，枚举左端点的时候，我们希望左端点按照$s$的单调上升顺序排序，于是想到了单调栈。具体算法是，开一个严格单调上升的栈（栈里存下标，但单调性是按照$s$在那个下标的值来排的），遍历$s$，当遍历到$s[i]$时，如果发现栈顶对应的数大于等于$s[i]$，则pop出栈顶，直到栈空或者栈顶对应的数小于$s[i]$为止，然后开始二分出最右边的满足$s[i]-s[x]\ge K$的那个$x$，然后更新答案（当然如果栈空了，则直接看$s[i]\ge K$是否成立，然后更新答案）。代码如下：

class Solution {
 public:
  int shortestSubarray(vector<int>& a, int k) {
    int n = a.size();
    vector<long> p(n + 1, 0);
    for (int i = 0; i < n; i++) p[i + 1] = p[i] + a[i];
    int res = n + 1;
    vector<int> stk(n + 1);
    int top = 0;
    for (int i = 0; i <= n; i++) {
      while (top && p[stk[top - 1]] >= p[i]) top--;
      if (top) {
        int l = 0, r = top - 1;
        while (l < r) {
          int mid = l + (r - l + 1 >> 1);
          if (p[i] - p[stk[mid]] >= k) l = mid;
          else r = mid - 1;
        }

        if (p[i] - p[stk[l]] >= k) res = min(res, i - stk[l]);
      }
      stk[top++] = i;
    }

    return res == n + 1 ? -1 : res;
  }
};

时间复杂度$O(n\log n)$，空间$O(n)$。

法2：单调队列。思路和上面差不多。维护一个从队头到队尾单调下降的队列（队列里存下标，但单调性按照$s$在那个下标的值来排的），然后遍历$s$，当遍历到$s[i]$的时候，先看一下队头与$i$之间的左开右闭区间的数字和是否大于等于$K$，如果是，则更新答案，并且队头左边的下标不可能成为答案，poll出队头；如果不是，那就看一下$s[i]$是否小于等于队尾，如果是则poll队尾，直到队空或者大于队尾，再将$i$入队。总而言之，队列里维护的是能作为区间左端点的下标。我们尝试用数学归纳法证明一下：
证明每次for循环开始的时候，队列里存放的是以$i$为右端点的最优解可能的区间的左端点（这个区间是左开右闭的，因为$s[i]-s[j]$是包含$A$的第$i$个数但不包含第$j$个数的，这里计数从$1$开始）。首次循环和第二次循环开始的时候，显然成立。假设第$k-1$次循环的时候也成立，此时队头作为左端点如果区间和大于等于了$K$，那答案的一个上界就是$i$减去队头，那么此时的队头当然不能成为第$k$次循环的最优解的左端点（因为长度更大了），所以应该出队，以此类推。接着，如果队尾对应的数大于等于$s[i]$，那么下次循环的时候如果最优解能以当前队尾为左端点，那更能以$i$为左端点，从而得到更优解，所以队尾应当出队，以此类推。由数学归纳法，结论成立。

代码如下：

class Solution {
 public:
  int shortestSubarray(vector<int>& a, int k) {
    int n = a.size();
    vector<long> p(n + 1, 0);
    for (int i = 0; i < n; i++) p[i + 1] = p[i] + a[i];
    deque<int> dq;
    int res = n + 1;
    for (int i = 0; i <= n; i++) {
      while (dq.size() && p[dq.back()] >= p[i]) dq.pop_back();
      while (dq.size() && p[dq.front()] <= p[i] - k) {
        res = min(res, i - dq.front());
        dq.pop_front();
      }
      dq.push_back(i);
    }

    return res == n + 1 ? -1 : res;
  }
};

时空复杂度$O(n)$。