Top K Elements: Mastering Priority Queues
Sorting an array to find the Kth largest element takes O(N log N) time. But what if the array has 10 million elements, and you only need the Top 3? Sorting everything is a massive waste of computation. Enter the Priority Queue.
Introduction
If an interview question contains the phrase "Top K", "Kth Largest", or "Kth Smallest", it is a massive flashing neon sign telling you to use a Heap (implemented in Java as a PriorityQueue).
While simply calling Arrays.sort(nums) and picking the Kth element works, it requires $O(N \log N)$ time. We can optimize this heavily by maintaining a strict, limited-size Heap.
The Min-Heap Strategy
To find the Kth Largest element, we actually use a Min-Heap. This sounds counter-intuitive at first, but the logic is brilliant.
We iterate through our massive array of numbers and blindly toss them into our Min-Heap. The Min-Heap automatically organizes itself so that the absolute smallest number in the pile is always sitting at the very top.
The trick is this: We never let the heap size exceed K.
Every time we add a number, we check if the size is greater than K. If it is, we immediately pop the top element off. Because it is a Min-Heap, popping the top element deletes the smallest number in our current pile.
The Result
By constantly throwing away the smallest numbers, by the time we finish iterating through our array, all the small numbers have been destroyed. The only things left surviving inside our heap of size K are the K largest numbers in the entire array.
Since it is a Min-Heap, the smallest of those K survivors is sitting right at the top. That number is exactly the Kth largest element!
Complexity Analysis
- Time Complexity: $O(N \log K)$ — We iterate through N elements. Adding or removing from a heap of size K takes $\log K$ time. This is drastically faster than $O(N \log N)$ when K is small compared to N.
- Space Complexity: $O(K)$ — We only ever store K elements in memory at any given time, regardless of how massive the input array is.
Implementation
class Solution {
public int findKthLargest(int[] nums, int k) {
// Create a Min-Heap
PriorityQueue<Integer> minHeap = new PriorityQueue<>();
for (int num : nums) {
minHeap.add(num);
// If the heap grows larger than K, remove the smallest element
if (minHeap.size() > k) {
minHeap.poll();
}
}
// The root of the heap is the Kth largest element
return minHeap.peek();
}
}