Solving Top K Problems with Heaps

1. Understanding the "Top K" Problem Pattern

Many interview questions boil down to finding a specific subset of elements from a larger collection: the top K scores, the K most frequent words, the K closest neighbors, and so on.

A simple, brute-force approach would be to sort the entire collection and then pick the first or last K elements. Given an array of size $N$, sorting takes $O(N\log N)$ time. While correct, this is often suboptimal. If you have a million items and only need to find the top 10, sorting all one million seems wasteful. This is where heaps provide a more efficient solution.

To begin, let's get a formal overview of the pattern and the different heap-based strategies.

2. The Optimal Strategy: Finding the Top K Largest Elements

As the article explained, the most efficient way to find the K largest elements is to use a min-heap of size K. This might seem counterintuitive, so let's break down why it works.

The min-heap will act as a container for the K largest elements found so far.

• The root of this min-heap will always be the smallest of these K elements—our "weakest contender."
• When we process a new element from the input array, we compare it to this root.
  • If the new element is smaller than or equal to the root, it cannot possibly be one of the top K largest, so we discard it.
  • If the new element is larger than the root, it deserves a spot in our top K group. We remove the root (the current weakest contender) and insert the new, larger element. The heap then rearranges itself to find the new weakest contender and place it at the root.

After iterating through all $N$ elements, our min-heap contains exactly the K largest elements from the entire array.

Complexity Analysis

• Time Complexity: We iterate through all $N$ elements. For each element, we perform a heap operation (a push and possibly a pop). Since the heap's size is capped at $K$, each operation takes $O(\log K)$ time. This gives a total time complexity of $O(N\log K)$.
• Space Complexity: The heap stores at most $K$ elements, so the auxiliary space complexity is $O(K)$.

This $O(N\log K)$ approach is a significant improvement over the $O(N\log N)$ sorting solution, especially when $K$ is much smaller than $N$.

3. Code Walkthrough: "Kth Largest Element in an Array"

Let's solidify this with a canonical interview problem: LeetCode 215, "Kth Largest Element in an Array." You're given an array nums and an integer k, and you need to return the kth largest element. For example, if nums = [3,2,1,5,6,4] and k = 2, the answer is 5.

The following video provides an excellent walkthrough of the optimal min-heap approach for this exact problem.

Here is a Python implementation of the logic discussed in the video. We use the heapq module, which is a min-heap implementation.

import heapq

def findKthLargest(nums, k):
    # Create a min-heap. We'll maintain its size at K.
    min_heap = []
    
    for num in nums:
        # If the heap has less than K elements, just add the current number.
        if len(min_heap) < k:
            heapq.heappush(min_heap, num)
        # If the current number is larger than the smallest element in the heap (the root)
        elif num > min_heap[0]:
            # heappushpop is an efficient way to push a new item and then 
            # pop and return the smallest item.
            # This is equivalent to a push followed by a pop.
            heapq.heappushpop(min_heap, num)
            
    # After iterating through all numbers, the root of the min-heap
    # is the Kth largest element.
    return min_heap[0]

# Example:
nums = [3, 2, 1, 5, 6, 4]
k = 2
print(f"The {k}nd largest element is: {findKthLargest(nums, k)}")  # Output: 5

4. Pattern Variations and The Symmetric Case

The "Top K" pattern is highly versatile. Understanding its core logic allows you to adapt it to various problem statements.

Finding the Top K Smallest Elements

What if you need to find the K smallest elements? The logic is perfectly symmetrical: you use a max-heap of size K.

• The max-heap stores the K smallest elements seen so far.
• The root is the "strongest" of these contenders (the largest of the K smallest).
• If you encounter a new element that is smaller than the root, you pop the root and push the new element.
• Since Python's heapq is a min-heap, you can simulate a max-heap by pushing the negative of each number.

Top K Frequent Elements and K Closest Points

The pattern extends beyond simple numerical values. The "priority" can be a frequency, a distance, or any other computed metric.

As the video shows:

• Top K Frequent Elements: This is a two-step process. First, use a hash map to count the frequencies of all elements ($O(N)$ time). Then, apply the "Top K Largest" pattern on the frequencies, using a min-heap of size K to find the K elements with the highest frequencies.
• K Closest Points to Origin: Here, you want the K elements with the smallest distance. This is a "Top K Smallest" problem. You use a max-heap of size K to keep track of the K points with the smallest distances seen so far. The priority is the distance from the origin. A common optimization is to use the squared distance to avoid floating-point calculations with sqrt.

Conclusion

You have now mastered a fundamental and highly practical interview pattern. The ability to recognize a "Top K" problem and immediately jump to the optimal $O(N\log K)$ heap solution demonstrates a strong grasp of data structures and algorithmic trade-offs.

Key Takeaways:

• "Top K" problems are about finding a small subset of elements from a large collection based on some ordering.
• The optimal solution uses a heap of size K, achieving $O(N\log K)$ time and $O(K)$ space complexity.
• To find the K largest elements, use a min-heap of size K.
• To find the K smallest elements, use a max-heap of size K.
• The pattern can be adapted to various "priorities," such as frequency or distance, by first computing that metric.

In our next lesson, we will explore another classic problem that is elegantly solved with heaps: Merging K Sorted Lists. This will further solidify your understanding of how heaps can manage and order elements from multiple sources efficiently.