OJ: LeetCode 862 Shortest Subarray with Sum at Least K - 单调队列

题目

862. Shortest Subarray with Sum at Least K

Given an integer array nums and an integer k, return the length of the shortest non-empty subarray of nums with a sum of at least k. If there is no such subarray, return -1.

A subarray is a contiguous part of an array.

Example 1:

Input: nums = [1], k = 1
Output: 1

Example 2:

Input: nums = [1,2], k = 4
Output: -1

Example 3:

Input: nums = [2,-1,2], k = 3
Output: 3

Constraints:

• 1 <= nums.length <= 105
• -105 <= nums[i] <= 105
• 1 <= k <= 109

题解

知识点：单调队列

我们用数组p表示数组nums的前缀和，p[i]=sum(nums[:i]), 0<=i<=len(nums). 我们需要找到head和tail使得p[tail]-p[head] >= k且tail-head最小。从head到tail的范围在下文记为span。

我们从前往后读数组，读到每一个位置时，这个位置一定是一个潜在的头（possible head），而其是否能作为尾，需要根据现有的头决定。所以我们需要一个数据结构保存潜在的头，读到每一个位置时，先判断这个位置是否能作为尾，和潜在的头组成span，更新最短长度；然后将这个位置加入潜在的头。我们把这个过程细化成下面两步：

1. 判断某位置是否能作为尾，更新最短长度

在潜在的头里面找到所有和此位置相差k以上的头，更新最短长度。另外，这些所有的头在这次判断后将会抛弃，因为后面的位置和这些头组成的span肯定不是最短的。

2. 将某位置加入潜在的头

在将此位置加入潜在的头时，需要抛弃所有在此位置之前，且前缀和不小于此位置前缀和的头。因为当这些头能和后面的位置组成span时，此位置也能和同样的后面位置组成span，且后者更短。

以上就是我们解决这道题的步骤。不难发现，以上的步骤需要“潜在的头”这个数据结构满足两条性质：

1. 任意给定值，能找到前缀和小于该值的所有头
2. 任意给定值，能找到前缀和不小于该值的所有头

所以很自然想到“潜在的头”是一个前缀和有序的双向队列，两条性质分别对应从双向队列的两头找值。以下记“队头”为前缀和较小的一端，“队尾”为前缀和较大的一端。

这样算法也就明朗了。对于每个位置遍历，进行以下操作：

1. 本位置作为尾：如果队列为空则跳过。否则，如果队头和本位置的前缀和之差不小于k，更新最短长度并弹出队头。重复此判断，至队列为空或前缀和之差小于k为止。
2. 本位置作为头：如果队列为空则直接加入。否则，如果队尾的前缀和不小于本位置，弹出队尾。重复此判断，至队列为空或队尾前缀和小于本位置为止。将本位置从队尾加入.

由于位置加入队列的顺序，队列始终保证从队头到队尾，前缀和单调递增，位置单调递增。这就是「单调队列」在这道题中的应用。

代码：

import collections

CONST_MAX = int(1e6)

class Solution:
    def shortestSubarray(self, nums: List[int], k: int) -> int:
        # prefix sum
        p = [0]
        for n in nums:
            p.append(p[-1]+n)
        # main process
        q = collections.deque() # possible heads
        shortest = CONST_MAX
        for i, n in enumerate(p):
            # could be tail of any possible head?
            while len(q) > 0 and n - p[q[0]] >= k:
                shortest = min(shortest, i - q[0])
                q.popleft()
            # remove impossible heads from q 
            while len(q) > 0 and n <= p[q[-1]]:
                q.pop()
            q.append(i)
        if shortest != CONST_MAX: return shortest
        else: return -1