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.

• 주어진 배열에 대해 이어지는 부분 배열의 합이 k가 되는 가장 짧은 부분 배열의 길이를 구하시오.

예시

[2,-1,2]

• 2 + -1 + 2 = 3 이 되어 길이는 3이다.

제한

• $1 <= nums.length <= 10^5$
• $10^5 <= nums[i] <= 10^5$
• $1 <= k <= 10^9$

풀이

• 슬라이딩 윈도우와 비슷하긴 한데, 음수가 들어갈수 있기 때문에 prefix_sum 이 오름차순이 아니게 된다.
• 이를 해결하기 위해 minHeap을 사용한다!
• 앞부분을 sliding window 대신 minheap으로 대체한다고 풀면 된다.

class Solution:
    def shortestSubarray(self, nums: List[int], k: int) -> int:
        heap = []

        N = len(nums)

        prefix = [0] * N

        ans = len(nums) + 1

        for right in range(N):
            if right == 0:
                prefix[right] = nums[right]
            else:
                prefix[right] = prefix[right - 1] + nums[right]

            if prefix[right] >= k:
                ans = min(ans, right + 1)

            while heap and prefix[right] - heap[0][0] >= k:
                ans = min(ans, right - heap[0][1])
                heappop(heap)

            heappush(heap, (prefix[right], right))

        return ans if ans <= len(nums) else -1

• 혹은 prefix의 값이 증가하게 하여 해결할수도 있다.
• Ex. [-1, 2, -1, ...] 이 있다고 할때
  • (A) [-1] prefix 는 합이 -1이고
  • (B) [-1, 2] prefix 는 합이 2이고
  • (C) [-1, 2, -1] prefix는 합이 1이 된다.
  • 그런데 B경우의 prefix가 빠지더라도 C보다 길이가 1개 더 길지만 오히려 값은 더 줄어들게 된다.
  • 즉 고려할 필요가 없는 경우가 되는 것이다!

class Solution:
    def shortestSubarray(self, nums: List[int], k: int) -> int:
        res = float('inf')

        cur = 0
        q = deque()

        for r in range(len(nums)):
            cur += nums[r]

            if cur >= k:
                res = min(res, r + 1)

            while q and cur - q[0][0] >= k:
                prefix, index = q.popleft()
                res = min(res, r - index)

            while q and q[-1][0] > cur:
                q.pop()

            q.append((cur, r))

        return -1 if res == float('inf') else res