K-way Merge

If you are preparing for technical coding interviews, you are likely familiar with the classic Merge Sort algorithm. Merge sort works by splitting a list into two halves, sorting them, and then merging the two sorted halves back together.

But, what happens when things get more complex? What if, instead of two sorted lists, you are given k sorted lists, e.g. 5, 10, or even 100, and asked to merge them all into one unified sorted list?

This scenario is a common and important interview pattern known as the K-Way Merge.

Understanding the pattern with a real-world example

Imagine you have k sorted lines of students, each arranged by their student ID (from lowest to highest). Your goal is to combine them into one single, master list that remains perfectly sorted.

Because every individual line is already sorted, you don’t need to search through every student at once. You only need to look at the first student in each of the k lines.

At each step, you compare the front student from all k lines and select the one with the lowest ID. That student is placed into the final line. Then, you move forward in the same line from which the student was taken, bringing the next student to the front for comparison.

You repeat this process until all students are placed in the final line. Each time, you choose the candidate with the lowest ID among the k current candidates and replace it with the next from the same line. This efficient merging method is called the K-Way Merge.

How to implement K-Way Merge

The secret behind an efficient K-Way Merge is the min-heap (also known as a Priority Queue). The heap allows us to constantly track the smallest available value across all k lists. It acts as a waiting room that constantly keeps the smallest available value from all k lists right at the top.

To implement this efficiently, the algorithm follows these structured steps:

1. Initialize the heap: Take the first element from each of the k sorted lists. Place these elements into a min-heap.
2. Track metadata: For every element in the heap, store its value, the index of the list it came from, and its position within that list.
3. Extract the minimum: Remove the smallest element from the top of the min-heap and add it to your final result list.
4. Replenish the heap: Immediately look at the list that provided the element you just removed. If that list has a next element, push it into the min-heap.
5. Iterate: Repeat this process of extracting and replenishing until the min-heap is empty.

Now, let’s look at the following illustration to get a better understanding of K-Way Merge algorithm.

1 / 7

Python implementation

The following Python implementation provides a clean and interview-ready implementation of K-Way Merge using a min-heap. By utilizing the heapq module, we efficiently identify the smallest element across all k lists.

import heapq

def k_way_merge(sorted_lists):
    min_heap = []
    result = []

    # Step 1: Initialize heap with the first element of each non-empty list
    # Heap item format: (value, list_index, element_index)
    for list_index, arr in enumerate(sorted_lists):
        if arr:  # ignore empty lists
            heapq.heappush(min_heap, (arr[0], list_index, 0))

    # Step 2: Extract-min and replenish from the same list until heap is empty
    while min_heap:
        value, list_index, element_index = heapq.heappop(min_heap)  # smallest available
        result.append(value)  # add to final merged result

        # Step 3: Push the next element from the same list (if it exists)
        next_index = element_index + 1
        if next_index < len(sorted_lists[list_index]):
            next_value = sorted_lists[list_index][next_index]
            heapq.heappush(min_heap, (next_value, list_index, next_index))

    return result


# ------------------ Example Input ------------------
if __name__ == "__main__":
    lists = [
        [1, 4, 9],
        [2, 3, 8],
        [0, 7, 10]
    ]

    print("Input (k sorted lists):")
    for i, lst in enumerate(lists):
        print(f"  List {i}: {lst}")

    merged = k_way_merge(lists)

    print("\nMerged Output:")
    print(merged)

Time complexity

The time complexity of the algorithm above is O(n log k), where n is the total number of elements across all lists and k is the number of lists. This is because each element is inserted into the heap once and removed once, and each heap operation takes O(log k) as the heap never holds more than k elements at a time.

Space complexity

The space complexity of the algorithm is O(k) in auxiliary space because the heap stores at most one active element from each list. If you also count the output array that stores the merged result, the total space becomes O(n + k), which is effectively O(n) since the output dominates.

Common mistakes to avoid

While the core idea of K-Way Merge is straightforward, several implementation mistakes frequently appear in interviews. Being aware of these pitfalls can help you avoid unnecessary bugs and performance issues.

1. Not handling empty lists during heap initialization.
2. Pushing all elements from all lists into the heap upfront, which degrades runtime to O(n log n) instead of O(n log k).
3. Forgetting to push the next element from the same list after removing one from the heap.
4. Skipping bounds checks before accessing or pushing the next element.
5. Using an incorrect heap comparator (or accidentally using a max-heap instead of a min-heap).
6. Failing to store sufficient metadata in the heap (such as list index and element index).
7. Mishandling duplicate values due to flawed comparison logic.
8. In linked-list problems, losing node references or incorrectly connecting nodes, which can break the final list.