862. 和至少为 K 的最短子数组
题目描述
输入整数数组 nums 和整数 k,nums 中可能有负数。
子数组指数组中连续且非空的一段。找出元素和不小于 k 的子数组中长度最短的那个,返回其长度。
如果不存在这样的子数组,返回 -1。
输入:nums = [1], k = 1 输出:1
解释:[1] 本身的和就等于 1。
输入:nums = [1,2], k = 4 输出:-1
解释:整个数组和为 3,达不到 4。
输入:nums = [2,-1,2], k = 3 输出:3
解释:只有整个数组的和 2-1+2=3 达到 3。
数据范围
1 <= nums.length <= 10^5-10^5 <= nums[i] <= 10^51 <= k <= 10^9
解题思路:前缀和 + 单调队列
用前缀和 pre,子数组 nums[i..j-1] 的和是 pre[j] - pre[i]。从左到右枚举右端 j,维护一个下标双端队列,其中 pre 值严格递增。① 队头 i 若满足 pre[j] - pre[i] ≥ k,就用 j - i 更新答案并弹出:以后更大的 j 只会更长,i 没用了。② 队尾若 pre[队尾] ≥ pre[j],弹出:以后用 j 作左端更短且和更大。③ 把 j 放入队尾。
为什么正确
队列中下标递增、pre 值严格递增;每个可能成为最优左端点的下标都还在队列里。
复杂度分析
时间 O(n),每个下标最多进出队列一次;空间 O(n)。
易错点
有负数时普通滑动窗口失效(扩大窗口和可能变小)。前缀和可能超出 32 位范围,C++/Java 要用 64 位整数。
参考代码
Python、Java、C++、Go、JavaScript、TypeScript 六种语言,逻辑一致。
class Solution:
def shortestSubarray(self, nums: List[int], k: int) -> int:
n = len(nums)
pre = [0] * (n + 1)
for i, x in enumerate(nums):
pre[i + 1] = pre[i] + x
ans = n + 1
dq = deque()
for j in range(n + 1):
while dq and pre[j] - pre[dq[0]] >= k:
ans = min(ans, j - dq.popleft())
while dq and pre[dq[-1]] >= pre[j]:
dq.pop()
dq.append(j)
return ans if ans <= n else -1class Solution {
public int shortestSubarray(int[] nums, int k) {
int n = nums.length;
long[] pre = new long[n + 1];
for (int i = 0; i < n; i++) {
pre[i + 1] = pre[i] + nums[i];
}
int ans = n + 1;
Deque<Integer> dq = new ArrayDeque<>();
for (int j = 0; j <= n; j++) {
while (!dq.isEmpty() && pre[j] - pre[dq.peekFirst()] >= k) {
ans = Math.min(ans, j - dq.pollFirst());
}
while (!dq.isEmpty() && pre[dq.peekLast()] >= pre[j]) {
dq.pollLast();
}
dq.addLast(j);
}
return ans <= n ? ans : -1;
}
}class Solution {
public:
int shortestSubarray(vector<int>& nums, int k) {
int n = nums.size();
vector<long long> pre(n + 1, 0);
for (int i = 0; i < n; i++) {
pre[i + 1] = pre[i] + nums[i];
}
int ans = n + 1;
deque<int> dq;
for (int j = 0; j <= n; j++) {
while (!dq.empty() && pre[j] - pre[dq.front()] >= k) {
ans = min(ans, j - dq.front()); dq.pop_front();
}
while (!dq.empty() && pre[dq.back()] >= pre[j]) {
dq.pop_back();
}
dq.push_back(j);
}
return ans <= n ? ans : -1;
}
};func shortestSubarray(nums []int, k int) int {
n := len(nums)
pre := make([]int, n+1)
for i, x := range nums {
pre[i+1] = pre[i] + x
}
ans := n + 1
dq := []int{}
for j := 0; j <= n; j++ {
for len(dq) > 0 && pre[j]-pre[dq[0]] >= k {
if j-dq[0] < ans { ans = j - dq[0] }; dq = dq[1:]
}
for len(dq) > 0 && pre[dq[len(dq)-1]] >= pre[j] {
dq = dq[:len(dq)-1]
}
dq = append(dq, j)
}
if ans <= n {
return ans
}
return -1
}var shortestSubarray = function(nums, k) {
const n = nums.length;
const pre = new Array(n + 1).fill(0);
for (let i = 0; i < n; i++) {
pre[i + 1] = pre[i] + nums[i];
}
let ans = n + 1;
const dq = new Array(n + 1); let head = 0, tail = 0;
for (let j = 0; j <= n; j++) {
while (head < tail && pre[j] - pre[dq[head]] >= k) {
ans = Math.min(ans, j - dq[head++]);
}
while (head < tail && pre[dq[tail - 1]] >= pre[j]) {
tail--;
}
dq[tail++] = j;
}
return ans <= n ? ans : -1;
};function shortestSubarray(nums: number[], k: number): number {
const n = nums.length;
const pre: number[] = new Array(n + 1).fill(0);
for (let i = 0; i < n; i++) {
pre[i + 1] = pre[i] + nums[i];
}
let ans = n + 1;
const dq: number[] = new Array(n + 1); let head = 0, tail = 0;
for (let j = 0; j <= n; j++) {
while (head < tail && pre[j] - pre[dq[head]] >= k) {
ans = Math.min(ans, j - dq[head++]);
}
while (head < tail && pre[dq[tail - 1]] >= pre[j]) {
tail--;
}
dq[tail++] = j;
}
return ans <= n ? ans : -1;
}其他解法
两者都枚举右端点并丢弃用过的左端点;单调队列另外淘汰前缀和不更小的旧下标,把堆的 O(log n) 降到均摊 O(1)。
解法 2:前缀和 + 小根堆 时间 O(n log n),空间 O(n)
枚举右端点 j,把之前的 (pre[i], i) 放进按前缀和排序的小根堆。只要堆顶满足 pre[j] - pre[i] ≥ k,就用 j - i 更新答案并弹出堆顶:以后的右端点再配这个 i 只会更长,所以可以丢掉。弹完后把 (pre[j], j) 压入堆。
只用到“用过的左端点可以丢弃”这一半观察,没有利用单调性淘汰受支配的左端点,所以多一个 log n。比单调队列更容易想到,是通往最优解的好台阶。
class Solution:
def shortestSubarray(self, nums: List[int], k: int) -> int:
n = len(nums)
ans = n + 1
heap = []
pre = 0
heapq.heappush(heap, (0, 0))
for j, x in enumerate(nums, 1):
pre += x
while heap and pre - heap[0][0] >= k:
ans = min(ans, j - heapq.heappop(heap)[1])
heapq.heappush(heap, (pre, j))
return ans if ans <= n else -1class Solution {
public int shortestSubarray(int[] nums, int k) {
int n = nums.length;
int ans = n + 1;
PriorityQueue<long[]> pq = new PriorityQueue<>((a, b) -> Long.compare(a[0], b[0]));
long pre = 0;
pq.offer(new long[]{0, 0});
for (int j = 1; j <= n; j++) {
pre += nums[j - 1];
while (!pq.isEmpty() && pre - pq.peek()[0] >= k) {
ans = Math.min(ans, j - (int) pq.poll()[1]);
}
pq.offer(new long[]{pre, j});
}
return ans <= n ? ans : -1;
}
}class Solution {
public:
int shortestSubarray(vector<int>& nums, int k) {
int n = nums.size();
int ans = n + 1;
priority_queue<pair<long long, int>, vector<pair<long long, int>>, greater<pair<long long, int>>> pq;
long long pre = 0;
pq.push({0, 0});
for (int j = 1; j <= n; j++) {
pre += nums[j - 1];
while (!pq.empty() && pre - pq.top().first >= k) {
ans = min(ans, j - pq.top().second);
pq.pop();
}
pq.push({pre, j});
}
return ans <= n ? ans : -1;
}
};type prefixHeap [][2]int
func (h prefixHeap) Len() int { return len(h) }
func (h prefixHeap) Less(i, j int) bool { return h[i][0] < h[j][0] }
func (h prefixHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *prefixHeap) Push(x interface{}) { *h = append(*h, x.([2]int)) }
func (h *prefixHeap) Pop() interface{} {
old := *h
x := old[len(old)-1]
*h = old[:len(old)-1]
return x
}
func shortestSubarray(nums []int, k int) int {
n := len(nums)
ans := n + 1
h := &prefixHeap{{0, 0}}
pre := 0
for j := 1; j <= n; j++ {
pre += nums[j-1]
for h.Len() > 0 && pre-(*h)[0][0] >= k {
top := heap.Pop(h).([2]int)
ans = min(ans, j-top[1])
}
heap.Push(h, [2]int{pre, j})
}
if ans <= n {
return ans
}
return -1
}var shortestSubarray = function(nums, k) {
const n = nums.length;
let ans = n + 1;
const heap = [];
const push = (item) => {
heap.push(item);
let i = heap.length - 1;
while (i > 0) {
const p = (i - 1) >> 1;
if (heap[p][0] <= heap[i][0]) break;
[heap[p], heap[i]] = [heap[i], heap[p]];
i = p;
}
};
const pop = () => {
const top = heap[0], last = heap.pop();
if (heap.length > 0) {
heap[0] = last;
let i = 0;
while (true) {
let s = i;
const l = 2 * i + 1, r = 2 * i + 2;
if (l < heap.length && heap[l][0] < heap[s][0]) s = l;
if (r < heap.length && heap[r][0] < heap[s][0]) s = r;
if (s === i) break;
[heap[s], heap[i]] = [heap[i], heap[s]];
i = s;
}
}
return top;
};
push([0, 0]);
let pre = 0;
for (let j = 1; j <= n; j++) {
pre += nums[j - 1];
while (heap.length > 0 && pre - heap[0][0] >= k) ans = Math.min(ans, j - pop()[1]);
push([pre, j]);
}
return ans <= n ? ans : -1;
};function shortestSubarray(nums: number[], k: number): number {
const n = nums.length;
let ans = n + 1;
const heap: number[][] = [];
const push = (item: number[]): void => {
heap.push(item);
let i = heap.length - 1;
while (i > 0) {
const p = (i - 1) >> 1;
if (heap[p][0] <= heap[i][0]) break;
[heap[p], heap[i]] = [heap[i], heap[p]];
i = p;
}
};
const pop = (): number[] => {
const top = heap[0], last = heap.pop()!;
if (heap.length > 0) {
heap[0] = last;
let i = 0;
while (true) {
let s = i;
const l = 2 * i + 1, r = 2 * i + 2;
if (l < heap.length && heap[l][0] < heap[s][0]) s = l;
if (r < heap.length && heap[r][0] < heap[s][0]) s = r;
if (s === i) break;
[heap[s], heap[i]] = [heap[i], heap[s]];
i = s;
}
}
return top;
};
push([0, 0]);
let pre = 0;
for (let j = 1; j <= n; j++) {
pre += nums[j - 1];
while (heap.length > 0 && pre - heap[0][0] >= k) ans = Math.min(ans, j - pop()[1]);
push([pre, j]);
}
return ans <= n ? ans : -1;
}