https://leetcode.com/problems/shortest-subarray-with-sum-at-least-k/
给定一个数组 A A A,求其最短的和大于等于 K K K的子数组的长度。题目保证 K > 0 K>0 K>0。
法1:单调栈 + 二分。首先要求子数组的和,想到前缀和。设数组 s s s是 A A A的前缀和,具体来说,规定 s [ 0 ] = 0 s[0]=0 s[0]=0, s [ i ] = A [ 0 ] + . . . + A [ i − 1 ] s[i]=A[0]+...+A[i-1] s[i]=A[0]+...+A[i−1]。那么问题转为求在 s [ y ] − s [ x ] ≥ K s[y]-s[x]\ge K s[y]−s[x]≥K的情况下的最小的 y − x y-x y−x。首先,在固定子数组右端点 i i i的情况下,考虑左端点,如果 x < y x
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\log n) O(nlogn),空间 O ( n ) O(n) O(n)。
法2:单调队列。思路和上面差不多。维护一个从队头到队尾单调下降的队列(队列里存下标,但单调性按照 s s s在那个下标的值来排的),然后遍历 s s s,当遍历到 s [ i ] s[i] s[i]的时候,先看一下队头与 i i i之间的左开右闭区间的数字和是否大于等于 K K K,如果是,则更新答案,并且队头左边的下标不可能成为答案,poll出队头;如果不是,那就看一下 s [ i ] s[i] s[i]是否小于等于队尾,如果是则poll队尾,直到队空或者大于队尾,再将 i i i入队。总而言之,队列里维护的是能作为区间左端点的下标。我们尝试用数学归纳法证明一下:
证明每次for循环开始的时候,队列里存放的是以 i i i为右端点的最优解可能的区间的左端点(这个区间是左开右闭的,因为 s [ i ] − s [ j ] s[i]-s[j] s[i]−s[j]是包含 A A A的第 i i i个数但不包含第 j j j个数的,这里计数从 1 1 1开始)。首次循环和第二次循环开始的时候,显然成立。假设第 k − 1 k-1 k−1次循环的时候也成立,此时队头作为左端点如果区间和大于等于了 K K K,那答案的一个上界就是 i i i减去队头,那么此时的队头当然不能成为第 k k k次循环的最优解的左端点(因为长度更大了),所以应该出队,以此类推。接着,如果队尾对应的数大于等于 s [ i ] s[i] s[i],那么下次循环的时候如果最优解能以当前队尾为左端点,那更能以 i 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 ) O(n) O(n)。