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.

Time O(log k) per addSpace O(k)Code C#
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

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.