和至少为 K 的最短子数组

标签: 队列 数组 二分查找 前缀和 滑动窗口 单调队列 堆(优先队列)

难度: Hard

给你一个整数数组 nums 和一个整数 k ,找出 nums 中和至少为 k最短非空子数组 ,并返回该子数组的长度。如果不存在这样的 子数组 ,返回 -1

子数组 是数组中 连续 的一部分。

示例 1:

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

示例 2:

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

示例 3:

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

提示:

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

Submission

运行时间: 204 ms

内存: 31.6 MB

class Solution:
    def shortestSubarray(self, nums: List[int], k: int) -> int:
        # contiguous: means NOT sort
        from itertools import accumulate
        # if sum(nums) < k:
        #     return -1
        stack = deque() # place index
        # preSum
        preSum_lst = list(accumulate(nums, initial = 0))
        res = float("inf")
        for i,v in enumerate(preSum_lst):
            # keep mono-increasing, assume keep larger than me;
            #       if find one smaller than me, check
          
            while stack and v < preSum_lst[stack[-1]]:
                stack.pop()
            while stack and v - preSum_lst[stack[0]] >= k:
                res = min(res, i - stack.popleft())
            stack.append(i)
        return -1 if res == float("inf") else res
            
            

Explain

此题解采用前缀和和单调队列的组合方法来解决问题。首先,使用前缀和数组来记录从数组开始到当前位置的所有元素的累加和,这样可以快速计算任意子数组的和。然后,使用一个单调递增的队列来存储前缀和数组的索引,队列中的每个元素代表一个潜在的最小起点位置。对于每个新的前缀和,我们首先从队尾开始移除那些不再可能成为最短子数组起点的元素(因为一个更小的前缀和已出现)。然后,从队首开始检查是否存在有效的子数组,即当前前缀和与队首索引对应的前缀和之差至少为k。如果存在,更新结果并弹出队首元素,因为更短的有效子数组不可能再以此队首元素为起点。最后,将当前前缀和的索引加入队列。这一过程持续到遍历完所有的前缀和,最后根据结果变量判断是否找到了满足条件的子数组。

时间复杂度: O(n)

空间复杂度: O(n)

class Solution:
    def shortestSubarray(self, nums: List[int], k: int) -> int:
        from itertools import accumulate
        from collections import deque
        # 前缀和数组,初始值为0以方便计算从头开始的子数组和
        preSum_lst = list(accumulate(nums, initial=0))
        res = float('inf')  # 结果初始化为无穷大
        stack = deque()  # 单调队列,存放前缀和的索引
        for i, v in enumerate(preSum_lst):
            # 维持队列的单调递增性
            while stack and v < preSum_lst[stack[-1]]:
                stack.pop()
            # 检查当前前缀和与队首前缀和之差是否满足条件
            while stack and v - preSum_lst[stack[0]] >= k:
                res = min(res, i - stack.popleft())
            stack.append(i)  # 将当前索引加入队列
        # 如果res仍为无穷大,说明没有找到满足条件的子数组
        return -1 if res == float('inf') else res

Explore

在前缀和数组中加入初始值0的主要好处是方便计算从数组第一个元素开始到当前元素的子数组和。有了初始值0,前缀和数组中的第i个元素(preSum_lst[i])就代表了原数组中从第一个元素到第i-1个元素的总和。这样,如果需要计算从数组开头到某个位置i的子数组和,可以直接用preSum_lst[i+1]表示,无需额外计算。此外,这也使得在使用单调队列来查找满足条件的子数组时,可以直接用当前前缀和与队列中的前缀和相减,以检查是否满足和至少为k的条件,从而简化了逻辑和实现。

在单调队列中移除队尾元素的目的是保持队列的单调递增特性。这样做的原因是,如果当前的前缀和小于队尾的前缀和,那么在检查后续元素时,使用当前的前缀和作为潜在的起点会更有可能找到满足条件的子数组且子数组长度更短,因为它提供了一个更低的起始和。保持队列单调递增确保了我们总是有可能找到最短的满足条件的子数组,从而优化了算法的效率和结果。

在找到一个满足条件的子数组后,从队列中弹出队首元素是为了避免重复计算以及确保找到的子数组是最短的。当队首元素与当前前缀和之差已经满足条件时,以该队首元素为起点的最短子数组已经被找到。保留该队首元素并继续检查可能会导致找到更长的符合条件的子数组,而我们的目标是找到长度最短的子数组。因此,一旦找到满足条件的子数组,就应该移除当前队首元素,以便检查是否有其他更短的子数组满足条件。

输入数组中包含负数会增加算法处理的复杂度,因为负数可以减少前缀和的值,从而影响子数组和的计算。有负数存在时,前缀和不再是严格递增的,这就需要算法在维持单调队列的过程中更频繁地做出调整,比如移除队尾的较大前缀和以保持队列的单调性。此外,负数的存在可能导致在更多的位置需要检查是否存在满足条件的子数组,因为前缀和的减少可能使得之前不满足条件的子数组在包含了负数后满足条件,这增加了算法的执行步骤和时间复杂度。