Skip to main content
Back to path

Welcome to a Grasp lesson

Create a free account to keep studying and create your own study path.

Lesson 40

Solving Top-K Problems with Min-Heaps

Recommended resources

Hello! Welcome back to our course on Python for competitive programming.

In our last lesson, we learned how to implement a priority queue using Python's heapq module. We saw that by storing items as tuples (priority, data), we can efficiently retrieve the item with the highest priority (the smallest priority number).

Today, we'll put that knowledge to immediate practical use. We're going to tackle a classic and very common category of problems: Top-K element problems. By the end of this lesson, you will be able to solve these problems efficiently using a min-heap. This is a fundamental pattern that frequently appears in coding interviews and competitive programming challenges.

1. The "Top-K" Problem

What is a "Top-K" problem? In general, it asks you to find the k largest, k smallest, or k most frequent elements from a large collection of data.

A simple example: Given the array [3, 2, 1, 5, 6, 4] and k = 2, the task is to find the 2nd largest element. If we sort the array to [1, 2, 3, 4, 5, 6], the largest is 6, and the second largest is 5.

The most obvious approach is to sort the entire array and then pick the k-th element from the end. This works, but it has a time complexity of O(N log N), where N is the total number of elements. When N is very large (e.g., millions of items) and k is small (e.g., 100), sorting the entire dataset is wasteful. We can do much better using heaps.

2. Finding the Kth Largest Element with a Heap

Let's focus on the problem of finding the Kth largest element. A heap is the perfect tool for this. To understand the optimal approach, it's helpful to first see a more intuitive, but slightly less efficient, heap-based solution.

Kth Largest Element in an Array - Leetcode 215 - Heaps (Python)

The following video from Greg Hogg provides an excellent walkthrough of this problem. First, watch the segment where he introduces the problem and explains a solution using a max-heap.

Watch from the beginning to 06:02. Pay attention to how a max-heap is simulated in Python (by negating numbers) and how it's used to find the Kth largest element by popping K times. Note the time complexity he discusses.

The max-heap approach works, but its time complexity is O(N + K log N). We can optimize this. The most efficient and standard solution uses a min-heap of a fixed size, k.

The Optimal Min-Heap Strategy

The core idea is to maintain a min-heap that stores only the k largest elements encountered so far.

Here’s how the algorithm works:

  1. Initialize an empty min-heap.
  2. Iterate through each number (num) in the input array.
  3. If the heap currently has fewer than k elements, simply push num onto the heap.
  4. If the heap already has k elements, compare num with the smallest element in the heap (which is always at the root, heap[0]).
    • If num is larger than the root of the min-heap, it means num belongs in our set of the top k elements. We remove the current smallest element (the root) and add num. The heapq.heappushpop() function does this efficiently.
    • If num is smaller than or equal to the root, we do nothing. This num cannot be one of the top k largest elements.

After iterating through the entire array, the min-heap will contain the k largest elements from the array. The root of the heap (heap[0]) will be the Kth largest element.

Finding k Largest Elements with Min-Heap
This image contrasts the simple sorting approach with the more efficient min-heap method. The min-heap, acting as a "gatekeeper" of size k, only keeps track of the k largest elements seen so far, avoiding a full sort.

This might seem counter-intuitive: why use a min-heap to find the largest elements? Because the min-heap's root gives us instant access to the smallest of the top K candidates. This allows us to quickly decide if a new, incoming element is worthy of being included in our top K set.

Now, let's see this superior min-heap strategy in action.

Kth Largest Element in an Array - Leetcode 215 - Heaps (Python)

Let's return to the video, which now explains this more optimized min-heap approach. The visualization is key to understanding why this works so well.

Watch from 06:02 to 11:26. Focus on the logic of maintaining a heap of size k and how elements are pushed and popped. Also, observe the final Python implementation.

Complexity Analysis:

  • Time Complexity: O(N log K). We iterate through N elements. For each element, we might perform a heap operation (heappush or heappushpop). Since the heap's size is capped at k, each operation takes O(log K) time. This is a significant improvement over O(N log N) when k is much smaller than N.
  • Space Complexity: O(K). We need extra space to store the k elements in the heap.
Test your understanding!

You are asked to find the 3rd largest element (k=3) in the array arr = [8, 5, 12, 9, 4, 17, 2]. Using the min-heap approach, trace the state of the heap as you process the array. What is the final state of the heap, and what is the 3rd largest element?

Show answer
  1. Heap: []. Process 8. Heap size < 3. Push 8. Heap: [8]
  2. Heap: [8]. Process 5. Heap size < 3. Push 5. Heap: [5, 8]
  3. Heap: [5, 8]. Process 12. Heap size < 3. Push 12. Heap: [5, 8, 12]
  4. Heap: [5, 8, 12]. Process 9. Heap size is 3. 9 > heap[0] (9 > 5). Pop 5, push 9. heappushpop does this. Heap: [8, 9, 12]
  5. Heap: [8, 9, 12]. Process 4. Heap size is 3. 4 < heap[0] (4 < 8). Do nothing. Heap: [8, 9, 12]
  6. Heap: [8, 9, 12]. Process 17. Heap size is 3. 17 > heap[0] (17 > 8). Pop 8, push 17. Heap: [9, 17, 12] (heap property is maintained)
  7. Heap: [9, 17, 12]. Process 2. Heap size is 3. 2 < heap[0] (2 < 9). Do nothing. Heap: [9, 17, 12]

After processing the whole array, the final heap contains [9, 12, 17] (in some heap-ordered form). The root of the min-heap is 9, which is the 3rd largest element.

For a concise written explanation and code, you can also refer to this article.

Find k largest elements in an array

This article from GeeksForGeeks provides a clear, text-based summary of the min-heap approach and a compact Python implementation. It's a good supplement to the video.

Read the section titled '[Expected Approach] Using Priority Queue(Min-Heap)'. Focus on the logic description and the Python code snippet provided.

3. A Common Variation: Top K Frequent Elements

The "Top-K" pattern is very flexible. Let's apply it to a slightly different problem: finding the k most frequent elements in an array.

Problem: Given nums = [1, 1, 1, 2, 2, 3] and k = 2, return the two most frequent elements. Here, 1 appears 3 times, 2 appears 2 times, and 3 appears once. The top 2 are 1 and 2.

The solution involves two main steps:

  1. Count Frequencies: First, we need to count how many times each number appears. Python's collections.Counter is perfect for this. It will give us a dictionary-like object: {1: 3, 2: 2, 3: 1}.
  2. Find Top K Frequencies: Now, the problem is transformed into finding the k items with the highest frequencies. This is the exact same "Top-K" pattern we just learned! We can use a min-heap of size k to keep track of the k elements with the highest frequencies seen so far. The items stored in our heap will be tuples of (frequency, element).

Top K Frequent Elements - Leetcode 347 - Heaps (Python)

This video, also from Greg Hogg, walks through solving the 'Top K Frequent Elements' problem, applying the same min-heap logic we've just discussed.

Watch from the beginning to 08:08. Pay close attention to how collections.Counter is used first, and then how the min-heap stores (frequency, element) tuples to solve the problem.

This variation shows how a more complex problem can be broken down into a familiar pattern. Recognizing when a problem is a "Top-K" problem in disguise is a key skill in competitive programming.

This heap-based approach is particularly powerful for processing streaming data, where you can't store the entire dataset. You can read more about the advantages and real-world applications of this technique.

Top K Elements: Mastering Heap and Quickselect Algorithms

To understand where this technique fits in the broader landscape of algorithms, this article compares the heap-based solution with another method called Quickselect and discusses real-world applications.

Skim through the sections 'Advantages of the Heap Approach', 'Real-world Applications', and the 'Conclusion'. This will give you a good perspective on why the heap method is so valuable, especially for streaming data.

Conclusion

You have now learned one of the most important and versatile applications of heaps. The "Top-K" pattern is a staple in algorithmic problem-solving.

Key Takeaways:

  • "Top-K" problems ask for the k largest/smallest/most frequent elements from a collection.
  • While sorting works (O(N log N)), it's often inefficient.
  • The optimal solution for finding the k largest elements is to use a min-heap of size k.
  • This approach has a time complexity of O(N log K) and a space complexity of O(K).
  • This pattern is flexible and can be adapted to solve variations, like finding the top k frequent elements, by first transforming the input (e.g., counting frequencies).

Up Next:
In the next lesson, we will explore another clever application of heaps: "Find the median in a data stream using two heaps". This problem builds on today's concepts and introduces a new way of combining heaps to solve even more complex challenges.