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:
- Iterate through the array and compute the cumulative sum at each index.
- Use a min-heap (priority queue) to store pairs of cumulative sum and their corresponding indices.
- 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.
- 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;
}