Kth Largest Element in a Stream

How it works

To maintain the k-th largest element in a stream of numbers, we do not need to store all values.
Instead, we only need to keep track of the k largest elements seen so far.

A min-heap of size k is perfect for this:

• A min-heap always keeps the smallest value at the top.
• If the heap contains the k largest elements, then the smallest among them is exactly the k-th largest overall.
• Whenever a new number arrives:
  • If we add it and the heap grows beyond k, we remove the smallest element - because it cannot be in the top k anymore.

This way, the heap always holds exactly the top k elements, and retrieving the k-th largest is O(1).

Steps

Initialization

1. Insert all initial numbers into a min-heap.
2. If the heap size becomes greater than k, repeatedly remove the smallest element.
   • After this, the heap contains exactly k elements.

add(value)

1. Insert the new value into the min-heap.
2. If heap size > k:
   • Remove the smallest element (the heap root).
3. Return the heap's smallest element (the root), which is now the k-th largest.

Complexity

• Time complexity: O(m * log k)
• Space complexity: O(k)

Where m is the number of calls made to add().

Things to watch for

Using a Max-Heap Instead of Min-Heap

A common mistake is using a max-heap to find the k-th largest element. While it seems intuitive to keep the largest elements at the top, a max-heap would require storing all elements and repeatedly extracting the maximum k times for each query. A min-heap of size k is the correct choice because the root always holds the k-th largest element directly, allowing O(1) retrieval after each insertion.

Forgetting to Maintain Heap Size

When using a min-heap, it is essential to remove the smallest element whenever the heap size exceeds k. Failing to do so means the heap grows unbounded, and the root no longer represents the k-th largest element. Always check and pop after each insertion to keep exactly k elements in the heap.

from typing import List
import heapq

class KthLargest:

    def __init__(self, k: int, nums: List[int]):
        self.minHeap, self.k = nums, k
        heapq.heapify(self.minHeap)
        while len(self.minHeap) > k:
            heapq.heappop(self.minHeap)

    def add(self, val: int) -> int:
        heapq.heappush(self.minHeap, val)
        if len(self.minHeap) > self.k:
            heapq.heappop(self.minHeap)
        return self.minHeap[0]

class KthLargest {
private:
    priority_queue<int, vector<int>, greater<int>> minHeap;
    int k;

public:
    KthLargest(int k, vector<int>& nums) {
        this->k = k;
        for (int num : nums) {
            minHeap.push(num);
            if (minHeap.size() > k) {
                minHeap.pop();
            }
        }
    }

    int add(int val) {
        minHeap.push(val);
        if (minHeap.size() > k) {
            minHeap.pop();
        }
        return minHeap.top();
    }
};