在问 · 算法题解

862. 和至少为 K 的最短子数组

困难二分查找前缀和 + 单调队列力扣原题 ↗

题目描述

输入整数数组 nums 和整数 k,nums 中可能有负数。

子数组指数组中连续且非空的一段。找出元素和不小于 k 的子数组中长度最短的那个,返回其长度。

如果不存在这样的子数组,返回 -1。

示例 1
输入:nums = [1], k = 1
输出:1

解释:[1] 本身的和就等于 1。

示例 2
输入:nums = [1,2], k = 4
输出:-1

解释:整个数组和为 3,达不到 4。

示例 3
输入:nums = [2,-1,2], k = 3
输出:3

解释:只有整个数组的和 2-1+2=3 达到 3。

数据范围

  • 1 <= nums.length <= 10^5
  • -10^5 <= nums[i] <= 10^5
  • 1 <= 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 -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 -1

二分查找 专题的其他题目

查看 二分查找 全部题目 →