Shortest Subarray with Sum at least K

XXX. Shortest Subarray with Sum at Least K

Sliding Window
Heap (Priority Queue)

Problem Statement:

Given an integer array nums and an integer k, find the length of the shortest contiguous subarray such that the sum is at least k. If there is no such subarray, return -1.

Algorithm:

  1. Iterate through the array and compute the cumulative sum at each index.
  2. Use a min-heap (priority queue) to store pairs of cumulative sum and their corresponding indices.
  3. For each index:
    • If the cumulative sum itself satisfies the condition, update the result.
    • Check the heap for the smallest prefix sum such that the difference between the current cumulative sum and the prefix sum satisfies the condition. Update the result if necessary.
    • Remove invalid heap entries since they are no longer relevant for forming a shorter subarray.
    • Add the current cumulative sum and index to the heap.
  4. Return the shortest subarray length if found, otherwise return -1.

Complexity:

Time: O(n log n) – Iterating through the array takes O(n), and each heap operation (insert or remove) is O(log n).
Space: O(n) – The heap stores at most n elements.

Java Implementation:


// Actually makes a lot of sense without too much thinking 
// Key is you want to check the heap for the smallest possible prefix sum 
// so you can maximize the value of K
// And you can throw away heap items cause at that point they are already valid and 
// you ain't finding a better solution with a greater index since length will only get longer...
public int shortestSubarray(int[] nums, int k) {
    int n = nums.length, shortestSubarrayLength = Integer.MAX_VALUE;
    long cumulativeSum = 0;
    // Min-heap to store cumulative sum and corresponding index
    PriorityQueue> prefixSumHeap = new PriorityQueue<>(
        (a, b) -> Long.compare(a.getKey(), b.getKey())
    );

    for (int i = 0; i < n; i++) {
        cumulativeSum += nums[i]; // Update cumulative sum

        // Check if current subarray from 0 to i satisfies the condition
        if (cumulativeSum >= k) 
            shortestSubarrayLength = Math.min(shortestSubarrayLength, i + 1);

        // Remove elements from the heap to form valid subarrays
        while (!prefixSumHeap.isEmpty() && cumulativeSum - prefixSumHeap.peek().getKey() >= k) 
            shortestSubarrayLength = Math.min(shortestSubarrayLength, i - prefixSumHeap.poll().getValue());

        // Add current cumulative sum and index to heap
        prefixSumHeap.offer(new Pair<>(cumulativeSum, i));
    }

    return shortestSubarrayLength == Integer.MAX_VALUE ? -1 : shortestSubarrayLength; // Return result
}
                

Python Implementation:


import collections

def shortestSubarray(nums, k):
    n = len(nums)
    cumulative_sum = 0
    shortest_subarray_length = float('inf')
    prefix_sum_queue = collections.deque()  # Stores (cumulative_sum, index)

    for i in range(n):
        cumulative_sum += nums[i]  # Update cumulative sum

        # Check if current subarray from 0 to i satisfies the condition
        if cumulative_sum >= k:
            shortest_subarray_length = min(shortest_subarray_length, i + 1)

        # Remove elements from the queue to form valid subarrays
        while prefix_sum_queue and cumulative_sum - prefix_sum_queue[0][0] >= k:
            shortest_subarray_length = min(shortest_subarray_length, i - prefix_sum_queue.popleft()[1])

        # Add current cumulative sum and index to the queue
        prefix_sum_queue.append((cumulative_sum, i))

    return shortest_subarray_length if shortest_subarray_length != float('inf') else -1
                

C++ Implementation:


#include 
#include 
#include 
#include 
using namespace std;

int shortestSubarray(vector& nums, int k) {
    int n = nums.size();
    long cumulative_sum = 0;
    int shortest_subarray_length = INT_MAX;
    deque> prefix_sum_queue; // Stores {cumulative_sum, index}

    for (int i = 0; i < n; ++i) {
        cumulative_sum += nums[i]; // Update cumulative sum

        // Check if current subarray from 0 to i satisfies the condition
        if (cumulative_sum >= k) 
            shortest_subarray_length = min(shortest_subarray_length, i + 1);

        // Remove elements from the queue to form valid subarrays
        while (!prefix_sum_queue.empty() && cumulative_sum - prefix_sum_queue.front().first >= k) {
            shortest_subarray_length = min(shortest_subarray_length, i - prefix_sum_queue.front().second);
            prefix_sum_queue.pop_front();
        }

        // Add current cumulative sum and index to the queue
        prefix_sum_queue.push_back({cumulative_sum, i});
    }

    return shortest_subarray_length == INT_MAX ? -1 : shortest_subarray_length;
}
                
Previous
Previous

Peak index in Mountain Array

Next
Next

Prune tree