Kth Largest Element in a Data Stream

Hello! In our previous lesson, we generalized your knowledge of iterative tree traversals to handle N-ary trees, showing how fundamental algorithms can adapt to different data structures. This capped off our exploration of tree traversal patterns.

Today, we pivot to a new and powerful data structure: the Heap. Unlike trees, which are often used for storing and traversing hierarchical data, heaps excel at managing and retrieving elements based on priority. We'll explore this by solving a classic interview problem: finding the k-th largest element in a data stream. This pattern is fundamental for problems involving "Top K" elements, which appear frequently in areas from system monitoring to financial data analysis.

By the end of this lesson, you will be able to apply a single heap (also known as a priority queue) to efficiently find and maintain the k-th largest element as new data arrives.

What is a Heap?

A heap is a specialized tree-based data structure that satisfies the heap property. It comes in two main flavors:

• Min-Heap: The value of each node is less than or equal to the value of its children. This means the smallest element in the entire heap is always at the root.
• Max-Heap: The value of each node is greater than or equal to the value of its children. This means the largest element is always at the root.

The most powerful feature of a heap is its efficiency: finding the minimum (or maximum) element is an $O(1)$ operation, while adding or removing elements takes $O(\log n)$ time, where $n$ is the number of elements in the heap. In most programming languages, heaps are readily available as PriorityQueue (Java, C++), heapq (Python), or similar libraries.

The Problem: K-th Largest Element in a Stream

Imagine you're tracking stock prices and need to know the 3rd highest price seen so far at any given moment. As new prices arrive, this 3rd highest price might change. The problem is to design a class that can handle an initial set of numbers and then efficiently return the k-th largest element after each new number is added.

A naive approach would be to add the new number to a list and sort the entire list every time. If the list has $N$ elements, this would take $O(N\log N)$ for each add operation, which is too slow for a real-time stream. We need a more efficient way to maintain just the information we need.

The Core Idea: A Min-Heap of Size K

This is where the heap pattern shines. The key insight is that we don't need to track all the numbers in the stream. To find the k-th largest element, we only need to keep track of the k largest elements seen so far.

This leads to a counter-intuitive but powerful choice: we will use a min-heap of size k. Why a min-heap, not a max-heap?

Let's think it through. Suppose k=3 and our heap stores the 3 largest numbers we've seen: {10, 8, 12}.

• A max-heap would look like [12, 10, 8]. The root is 12 (the 1st largest). To find the 3rd largest, we'd have to pop 3 times.
• A min-heap would look like [8, 10, 12]. The root is 8. Notice that this root is the smallest of the three largest elements—which is precisely the 3rd largest element overall.

By maintaining a min-heap of size k with the k largest elements, the k-th largest element is always at the root, available in $O(1)$ time.

The following resource from Algo.monster provides an excellent explanation of this intuition.

The Algorithm

Now that we have the core idea, let's formalize the algorithm for our KthLargest class, which will have a constructor and an add method.

1. Constructor: KthLargest(k, nums)

• Initialize an empty min-heap and store the value of k.
• Iterate through the initial list of numbers, nums. For each number, call our add logic (which we'll define next). This ensures the heap is correctly constructed from the start.

2. Method: add(val)

• Step 1: Add the new value val to the min-heap. This is an $O(\log k)$ operation.
• Step 2: Check the size of the heap. If it's now greater than k, it means we have k+1 elements. The smallest of these (the root) cannot be among the top k largest, so we remove it. This is also an $O(\log k)$ operation.
• Step 3: After these steps, the heap contains exactly the k largest elements seen so far. The root of the heap (heap[0]) is the k-th largest element, which we return.

The NeetCode video below provides a fantastic visual walkthrough of this process, from the initial brute-force idea to the final, optimized heap solution.

Implementation and Complexity

Here is a Python implementation based on the algorithm. Notice how the __init__ method simply calls the add method repeatedly, which is a clean way to handle the initial numbers.

import heapq

class KthLargest:
    def __init__(self, k: int, nums: list[int]):
        self.k = k
        self.min_heap = []
        for num in nums:
            self.add(num)

    def add(self, val: int) -> int:



        # If the heap isn't full yet, just add the element.
        if len(self.min_heap) < self.k:
            heapq.heappush(self.min_heap, val)



        # If the new value is larger than the smallest element in the heap
        # (the k-th largest so far), it deserves a spot.
        elif val > self.min_heap[0]:



            # heappushpop is an efficient way to push an item and then
            # pop the smallest item.
            heapq.heappushpop(self.min_heap, val)
        



        # The k-th largest element is at the root of our min-heap of size k.
        # The problem statement guarantees this will only be called when
        # there are at least k elements.
        return self.min_heap[0]

Note: This implementation is slightly more optimized. Instead of always pushing and then conditionally popping, it first checks if the new element is even a candidate to be in the top k. An even simpler implementation, as shown in the video, is to always push and then pop if the size exceeds k. Both are correct and have the same complexity.

Complexity Analysis

• Time Complexity:
  • Constructor: $O(N\log k)$, where $N$ is the number of initial elements in nums. This is because we perform $N$ add operations, each taking $O(\log k)$.
  • add method: $O(\log k)$. This is a huge improvement over the naive $O(N\log N)$ approach.
• Space Complexity: $O(k)$, as we only ever store k elements in our heap, regardless of how large the stream becomes.

Mastering the Pattern: Common Pitfalls

Since your goal is a deep and "neat" understanding, let's review common mistakes developers make with this pattern. Recognizing these will solidify your grasp of the concept.

Conclusion

In this lesson, you've been introduced to heaps and a foundational algorithm pattern for "Top K" problems. You now have a practical and efficient solution for finding the k-th largest element in a stream, a problem that models many real-world scenarios.

Key Takeaways:

• Heaps (Priority Queues) are the ideal data structure for problems requiring efficient retrieval of the minimum or maximum element.
• To find the k-th largest element, the most efficient method is to use a min-heap of size k.
• This heap stores the k largest elements seen so far, and its root is always the k-th largest element overall.
• This pattern provides an $O(\log k)$ time complexity for adding new elements and $O(k)$ space complexity, making it highly scalable for large data streams.

This is just the first step in leveraging heaps. In our next lesson, we will build directly on this concept to solve a more complex problem: finding the median of a data stream. This will require us to use not one, but two heaps working in tandem.