从Leetcode 每日一题练习继续讨论:
862. 和至少为 K 的最短子数组
862. Shortest Subarray with Sum at Least K
题解
本题求子数组的和满足大于等于K的全部子数组中长度最小的长度是多少。这种求子数组和的有关性质的题目首先可以想到使用前缀和来解决,但之前使用前缀和解题时是因为前缀和有一个很重要的性质即单调性,当所有数字均为非负数时前缀和是单调增的,因此我们可以根据题目要求来找到对应的下标,如同样是求子数组和至少为K的问题,若是单调增的前缀和,则根据下标为i对应的前缀和减去K的前缀和P来找到对应的小于等于P的前缀和对应的下标,即可知道所有子数组和至少为K的子数组。
但本题中存在负数,因此前缀和不是单调增的,这时可以考虑能否构造一个单调增的前缀和,则可使用单调栈来构造这样的前缀和,考虑栈顶前缀和的大小和新的前缀和的大小关系,若栈顶前缀和比新的前缀和大,则可弹出栈顶,因为对于还未访问到的前缀和,若后面的前缀和和栈顶的差大于等于k,则因为新的前缀和比栈顶小同时其对应的下标位置在栈顶的后面,则后面的前缀和和新前缀和的差必定也大于等于k且二者对应下标的差比和当前栈顶下标的差更小,故栈顶可以舍弃。
要求和不小于K的子数组就需要根据当前的前缀和计算出符合要求的前缀和大小,并在单调栈中找到不大于这个符合要求的前缀和大小的前缀和对应的下标,寻找这个符合要求的前缀和每次都要从头遍历单调栈,能否通过优化减少从头遍历的次数呢。很简单,题目要求找到满足要求的前缀和的最短长度,则我们只要找到刚好满足要求前缀和对应的下标位置,在这个下标之前的都可以直接丢弃,后续不再遍历这些位置。因为后面未遍历的前缀和要么比当前的大,要么比当前的小,比当前大则当前前缀和能取得的使子数组大于等于k的下标位置对于后面的同样是可以取得的,但后面到该下标的距离一定比当前下标大,如果比当前小,那么要取得使子数组大于等于k需要在当前前缀和对应的解的位置前面找更小的前缀和来使得子数组满足要求,这样得到的长度最多只能和当前前缀和对应的解的长度相同,不会更优。
根据这两个优化,可以构造一个单调的队列,使得两端都能弹出元素并满足上面的优化方法。先计算出前缀和,再不断遍历前缀和并求得满足条件的子数组长度与记录的最小长度比较并不断更新即可。
代码
class Solution {
public:
int shortestSubarray(vector<int>& nums, int k) {
int n = nums.size();
vector<long long> preSum(n + 1);
preSum[0] = 0;
for (int i = 0; i < n; i++) {
preSum[i + 1] = preSum[i] + nums[i];
}
int res = n + 1;
// 使用双端队列存储下标
deque<int> dq;
// 遍历前缀和数组
for (int i = 0; i <= n; i++) {
while (!dq.empty() && preSum[i] <= preSum[dq.back()]) {
dq.pop_back();
}
while (!dq.empty() && preSum[i] - preSum[dq.front()] >= k) {
res = min(res, i - dq.front());
dq.pop_front();
}
dq.push_back(i);
}
return res == n + 1 ? -1 : res;
}
};