Maintaining Top-K Items with a Bounded Heap

Hello! In our last lesson, we constructed a MinHeap class from scratch, giving us a powerful tool for efficiently managing ordered data. We saw how the add and poll operations work in $O(\log N)$ time by maintaining the heap's structural and ordering properties.

Today, we'll put that MinHeap to work and solve a classic interview problem. This lesson focuses on a fundamental algorithmic pattern: using a bounded heap to find the k largest or smallest items in a collection or from a continuous stream of data. Mastering this pattern is a significant step towards interview readiness, as it demonstrates a non-obvious and highly efficient use of a core data structure.

Our goal is to solve the "Kth Largest Element in a Stream" problem. This will not only solidify your understanding of heap operations but also show you how to choose the right tool to dramatically improve an algorithm's performance.

The Problem: Finding the Kth Largest Element

Imagine you're processing a real-time feed of stock prices, user scores, or sensor readings. A common requirement is to constantly know the k-th largest value seen so far. For example, what is the 10th highest price in the last hour?

Let's formalize this as an interview problem, "Kth Largest Element in a Stream" (similar to LeetCode #703):

• You are given an integer k and an initial stream of numbers.
• You need to implement a class that supports an add method.
• Each time add(val) is called, it should incorporate the new value val into the stream and return the current k-th largest element.

A simple, but inefficient, approach would be to store all the numbers in a list, add the new number, sort the list in descending order, and pick the element at index k-1. If you have N numbers, each add operation would take $O(N\log N)$ time. As the stream grows, this becomes prohibitively slow. We need a much smarter way.

The Bounded Min-Heap: An Elegant Solution

The key insight is that we don't need to store all the numbers. We only care about the k largest ones. Any number that isn't in this top k group is irrelevant to our answer.

This leads to a counter-intuitive but powerful idea: to find the k largest elements, we will use a min-heap of size k.

Why a min-heap?

• A min-heap gives us instant (O(1)) access to its smallest element (the root).
• If we maintain a min-heap that only contains the k largest elements seen so far, then the root of this heap is, by definition, the smallest of the largest—which is exactly the k-th largest element overall.

Let's make this concrete. If k=5 and our min-heap holds the 5 largest numbers seen so far, say {90, 95, 101, 105, 110}, the heap's root will be 90. This 90 is our 5th largest element. Any new number smaller than 90 can be ignored, as it won't change the top 5. If a new number like 103 arrives, it belongs in the top 5, and 90 gets kicked out. The min-heap is the perfect data structure to manage this "kicking out" process efficiently.

The following resource provides an excellent explanation of this intuition.

The Algorithm in Action

Now that we have the core idea, let's detail the algorithm for the add method, assuming we have our MinHeap class from the previous lesson.

1. Add the new element: Push the new value val into the min-heap.
2. Maintain the bound: If the heap's size is now greater than k, remove the smallest element by calling poll().
3. Return the answer: The k-th largest element is now at the root of the heap, so we can return it with peek().

This simple two-step process (add, then poll if oversized) ensures the heap never stores more than k elements and that those k elements are always the largest ones encountered.

The NeetCode video below walks through this exact logic. It uses Python, but the concepts are universal and map directly to our TypeScript implementation.

TypeScript Implementation

Let's translate this logic into a KthLargest class in TypeScript. This class will use the MinHeap we built in the previous lesson.

// Assuming the MinHeap class from the previous lesson is available
// class MinHeap {
//   public add(value: number): void { ... }
//   public poll(): number | null { ... }
//   public peek(): number | null { ... }
//   public size(): number { ... }
// }

class KthLargest {
    private k: number;
    private minHeap: MinHeap;

    constructor(k: number, nums: number[]) {
        this.k = k;
        this.minHeap = new MinHeap();

        // Process the initial stream of numbers
        for (const num of nums) {
            this.add(num);
        }
    }

    public add(val: number): number | null {
        // If the heap has less than k elements, it's not full yet.
        // Or, if the new value is larger than the smallest of the top k elements (the heap's root),
        // then this new value belongs in our set of top k elements.
        if (this.minHeap.size() < this.k || val > this.minHeap.peek()!) {
            this.minHeap.add(val);
        }

        // If adding the new element made the heap too large,
        // remove the smallest element to maintain the size k bound.
        if (this.minHeap.size() > this.k) {
            this.minHeap.poll();
        }

        // The root of the heap is the kth largest element.
        // The problem often guarantees we can return a valid number,
        // but it's good practice to handle the case where the heap isn't full yet.
        if (this.minHeap.size() < this.k) {
            return null; // Or handle as per problem spec
        }
        
        return this.minHeap.peek();
    }
}

A slightly more concise way to write the add method, which is often seen in solutions, is to always add the element and then immediately trim the heap if it exceeds size k.

    public add(val: number): number | null {
        // 1. Add the new value unconditionally.
        this.minHeap.add(val);

        // 2. If the heap is now larger than k, remove the smallest element.
        if (this.minHeap.size() > this.k) {
            this.minHeap.poll();
        }
        
        // 3. Return the root, which is the k-th largest element.
        if (this.minHeap.size() < this.k) {
            return null;
        }
        return this.minHeap.peek();
    }

Both implementations are correct and achieve the same result. The second one is just a bit more direct.

The Algo.monster resource provides a complete implementation using a library MinPriorityQueue, but the logic is identical to what we've written.

Complexity Analysis

Let's analyze the efficiency of our bounded heap approach.

• Space Complexity: We only ever store k elements in our heap, regardless of how many numbers we process. The space complexity is therefore $O(k)$.
• Time Complexity:
  • Constructor: We process the initial N numbers. For each number, we perform an add operation on a heap of size at most k. This takes $O(\log k)$ time. The total time for initialization is $O(N\log k)$.
  • add method: Each call involves one add ($O(\log k)$) and possibly one poll ($O(\log k)$). The total time per call is $O(\log k)$.

This is a huge improvement over the $O(N\log N)$ time per addition of the naive sorting approach.

Generalizing the Pattern

This bounded heap pattern is not limited to just one problem. It's a general technique for finding the "Top K" of anything. Consider another famous problem: "Top K Frequent Elements" (LeetCode #347).

The goal is to find the k most frequent elements in an array. The solution uses the same pattern:

1. Use a Map to count the frequency of each number.
2. Create a min-heap of size k that stores [element, frequency] pairs. The heap should be ordered by frequency.
3. Iterate through your frequency map. For each [element, frequency] pair, add it to the heap. If the heap size exceeds k, poll() the pair with the lowest frequency.
4. After iterating through all elements, the heap will contain the k most frequent elements.

This shows the power of recognizing patterns. Once you understand the bounded heap, you can apply it to a whole family of problems.

Conclusion

Today we explored the bounded heap pattern, a crucial technique for solving "Top K" style problems efficiently. This is a common pattern in coding interviews because it requires you to think beyond the obvious solution and apply the right data structure for the job.

Key Takeaways:

• Bounded Heap Pattern: To find the k largest (or smallest) items from a large collection, maintain a heap of size k.
• Min-Heap for K Largest: Use a min-heap of size k to track the k largest elements. The root of the heap is always the k-th largest element.
• Max-Heap for K Smallest: Conversely, you would use a max-heap of size k to find the k smallest elements.
• Efficiency: This pattern offers excellent performance, with $O(\log k)$ time per addition and $O(k)$ space, making it ideal for processing large streams of data.

We have now concluded our module on trees and heaps. In the next module, we will venture into an even more general and powerful data structure: graphs. We will start by learning how to represent relationships between entities as a graph and then explore fundamental traversal algorithms like Breadth-First Search (BFS) and Depth-First Search (DFS). The heap we've studied will make a return appearance in one of the most famous graph algorithms, Dijkstra's algorithm for finding shortest paths.