Heap-based Solutions for Top-K and Merge K-Sorted Lists

Hello! Welcome to your next lesson in the "Advanced Algorithmic Patterns" module.

In our last lesson, we saw how a monotonic deque provides an optimal O(N) solution for finding the single maximum element in a sliding window. This was powerful, but it's tailored to finding just one "best" item at a time. Today, we'll address a different, but equally common, class of problems: what if you need to find the "Top K" items from a large dataset, or merge results from k different sorted sources?

These scenarios are frequent in real-world systems. For example, a service like Bedrock might need to identify the top k most frequently used models to optimize caching, or merge sorted log streams from k different server instances. For these tasks, we need a data structure that can efficiently manage a collection of prioritized elements. This is the perfect job for a heap, also known as a priority queue.

By the end of this lesson, you will be able to apply heaps to solve two classic problems: finding the top-K frequent elements and merging k-sorted lists, both of which are staples in software engineering interviews.

---

1. What is a Heap?

A heap is a specialized tree-based data structure that satisfies the "heap property." This property ensures that the parent node is always ordered relative to its children nodes. There are two main types:

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

Because heaps make it extremely efficient to find the min/max element, they are the ideal data structure for implementing priority queues. In a priority queue, every element has a "priority," and the element with the highest (or lowest) priority is processed next. This is a concept you've likely worked with extensively when building the priority-based queuing for the Rudi inference scheduler.

To understand how heaps maintain this property during insertions and deletions, let's watch a short video that covers the core mechanics.

The key takeaways on performance are:

• Insertion (push): O(log n)
• Deletion of min/max (pop): O(log n)
• Accessing min/max (peek): O(1)

Python's standard library provides the heapq module, which implements a min-heap. We will use this for our implementations.

---

2. Solving the "Top K Frequent Elements" Problem

This is a very common interview question: given a list of items, find the k most frequent ones.

Example: nums = [1, 1, 1, 2, 2, 3], k = 2.
The frequencies are: 1: 3 times, 2: 2 times, 3: 1 time.
The top 2 most frequent are 1 and 2.

The solution is a two-step process:

1. Count Frequencies: Iterate through the input list and store the frequency of each item in a hash map. This is an O(N) operation. Python's collections.Counter is perfect for this.
2. Find the Top K: Now, from the frequency map, we need to find the k items with the highest counts. A naive approach would be to sort all the unique items by frequency, which would take O(M log M) time, where M is the number of unique items. We can do better with a heap.

The Heap-Based Approach

The most efficient method uses a min-heap of size k. The core idea is to maintain a heap that holds the "best k candidates so far."

Let's watch a video that explains this strategy.

As the video explains, the algorithm is:

1. Initialize an empty min-heap.
2. Iterate through the (item, frequency) pairs from your hash map.
3. For each pair, push it onto the heap.
4. If the heap's size exceeds k, pop the smallest element (the root of the min-heap).
5. After iterating through all items, the heap contains the top k elements.

This approach has a time complexity of $O(M\log k)$, where M is the number of unique items. Since we first iterate through the original list of N items to build the frequency map, the total complexity is $O(N+M\log k)$. This is significantly better than sorting when k is much smaller than M.

Implementation in Python

Here’s how you would implement this using Python's heapq. Since heapq is a min-heap, we'll store (frequency, item) tuples. The heap will automatically use the first element of the tuple (the frequency) for ordering.

import heapq
from collections import Counter

def top_k_frequent(nums, k):



    # Step 1: Count frequencies - O(N)
    if not nums:
        return []
    counts = Counter(nums) # e.g., {1: 3, 2: 2, 3: 1}




    # Step 2: Use a min-heap to find top k - O(M log k)
    # M is the number of unique elements
    min_heap = []
    for item, freq in counts.items():



        # Push a tuple (frequency, item)
        heapq.heappush(min_heap, (freq, item))
        



        # If heap size exceeds k, remove the smallest element
        if len(min_heap) > k:
            heapq.heappop(min_heap)




    # The heap now contains the top k elements.
    # Extract just the items into a result list.
    result = [item for freq, item in min_heap]
    return result




# Example
nums = [1, 1, 1, 2, 2, 3]
k = 2
print(top_k_frequent(nums, k)) # Output will be [2, 1] (order may vary)

Deeper Dive on Trade-offs

Your preference for understanding trade-offs is crucial for senior roles. The choice between sorting and using a heap is a classic example.

The article you just read also alludes to another heap strategy: build a max-heap of all M items (O(M)) and then pop k times (O(k log M)). This gives a total of $O(M+k\log M)$. For a typical interview scenario where k is small and M can be large, the O(M \log k) min-heap approach is usually preferred, especially for its memory benefits with streaming data.

---

3. Solving "Merge K Sorted Lists"

Another classic problem where heaps excel is merging k sorted lists into a single sorted list. This problem pattern appears when aggregating results from distributed workers, a scenario you're familiar with from your work at AWS. Imagine each of k nodes returns a sorted list of results; you need an efficient way to merge them on a central aggregator.

Algorithm:

1. Initialize a min-heap.
2. Push the first element from each of the k lists into the heap. To keep track of where each element came from, push a tuple: (value, list_index, element_index_within_list).
3. Initialize an empty result list.
4. While the heap is not empty:
   a. Pop the smallest element tuple from the heap. This is the next smallest element in the overall merged list. Add its value to your result list.
   b. Look at the list_index and element_index from the popped tuple. If the list it came from still has more elements, push the next element from that list into the heap.

This process ensures that you are always considering the smallest available element across all k lists.

Complexity Analysis:

• Time: $O(N\log k)$, where N is the total number of elements across all lists and k is the number of lists. Each of the N elements will be pushed onto and popped from the heap exactly once. The heap size never exceeds k.
• Space: $O(k)$ to store at most one element from each of the k lists in the heap.

Implementation in Python

import heapq

def merge_k_lists(lists):
    min_heap = []
    



    # Step 1 & 2: Push the first element of each list into the heap
    for i, lst in enumerate(lists):
        if lst: # Ensure the list is not empty



            # Push (value, list_index, element_index)
            heapq.heappush(min_heap, (lst[0], i, 0))

    result = []
    



    # Step 4: Process until the heap is empty
    while min_heap:



        # a. Pop the smallest element
        val, list_idx, elem_idx = heapq.heappop(min_heap)
        result.append(val)
        



        # b. If there's a next element in the same list, push it
        if elem_idx + 1 < len(lists[list_idx]):
            next_elem_idx = elem_idx + 1
            next_val = lists[list_idx][next_elem_idx]
            heapq.heappush(min_heap, (next_val, list_idx, next_elem_idx))
            
    return result




# Example
list1 = [1, 4, 5]
list2 = [1, 3, 4]
list3 = [2, 6]
print(merge_k_lists([list1, list2, list3])) 



# Output: [1, 1, 2, 3, 4, 4, 5, 6]

The "System Design Interview" video you watched earlier briefly mentions this exact pattern when discussing how to aggregate Top-K results from multiple distributed processors (see timestamps 09:12 and 33:49 in resource LINK). This reinforces its practical relevance in building scalable systems.

---

Conclusion

Today, we've added the heap to our toolkit, a data structure perfectly suited for problems involving priorities and ordering a subset of elements.

Key Takeaways:

• Heaps (Priority Queues): A heap is a tree-based structure that maintains the heap property (min or max). It offers O(log n) insertion/deletion and O(1) peek operations, making it ideal for priority management.
• Top K Frequent Elements: This is solved efficiently with a combination of a hash map for counting (O(N)) and a min-heap of size k to track the top candidates (O(M log k)).
• Merge K Sorted Lists: A min-heap of size k is used to efficiently find the next smallest element across all lists, leading to an optimal O(N log k) solution.
• Python's heapq: This module provides a min-heap implementation that is robust and easy to use for these patterns.

Next Lesson Preview

We've now seen two powerful heap-based solutions and briefly touched on their complexity compared to sorting. In the next lesson, we will formalize this by analyzing the time and space complexity trade-offs of heap-based solutions versus sorting-based alternatives in more detail. This will equip you to justify your algorithmic choices with precision, a key skill evaluated in senior engineering interviews.