Kth Largest Element in a Stream
Keep a min-heap of the k largest values seen so far. Its root, the smallest of them, is the k-th largest overall.
Watch it animatedStep through it, with the running line highlighted. Free with an account.The idea
Only the k largest values can ever matter for "the k-th largest", so keep exactly those in a min-heap. The heap's root is the smallest of the k largest, which is the k-th largest overall. A new value that is not bigger than the root cannot enter the top k and is dropped. A bigger one replaces the root, which falls out of the top k, and sifts down. Each add costs O(log k), however long the stream gets.
Code
public class KthLargest {
readonly PriorityQueue<int, int> heap = new(); // min-heap: the k largest so far
readonly int k;
public KthLargest(int k, int[] nums) {
this.k = k;
foreach (int x in nums) Add(x);
}
public int Add(int val) {
if (heap.Count < k) heap.Enqueue(val, val); // fewer than k kept: keep everything
else if (val > heap.Peek()) heap.DequeueEnqueue(val, val); // beats the k-th largest
// else val <= the k-th largest, so it can never be in the top k
return heap.Peek();
}
}Tips
- A min-heap for "largest" values: the root is the one pushed out when something bigger arrives, so it must be the cheapest to reach.
- Skipped values are gone for good, which is safe: once k bigger values exist, a smaller value can never become the k-th largest again.
- The constructor just calls Add for each initial number, so the heap holds at most k values from the start.
In the .NET library
PriorityQueue<TElement, TPriority> is a min-heap on the priority. DequeueEnqueue (.NET 7+) removes the root and inserts in one sift; EnqueueDequeue(val, val) does the compare-and-skip for you, returning val straight away when it is not bigger than the root.