Coding Patterns: K-way Merge

In Coding Patterns series, we will try to recognize common patterns underlying behind each algorithm question, using real examples from Leetcode.

Previous posts were about Sliding Window, Two Pointers, Fast & Slow Pointers, Merge Intervals, Cyclic Sort, In-place Reversal of a Linked List, Breadth First Search (BFS), Depth First Search (DFS), Two Heaps, Subsets, Modified Binary Search and Top K Numbers patterns and today, we will introduce K-way Merge pattern which is very useful to solve the problems whenever we are given K sorted arrays, we can use a Heap to efficiently perform a sorted traversal of all the elements of all arrays.

We can push the smallest (first) element of each sorted array in a Min Heap to get the overall minimum. While inserting elements to the Min Heap we keep track of which array the element came from.

We can also remove the top element from the heap to get the smallest element and push the next element from the same array, to which this smallest element belonged, to the heap. We can repeat this process to make a sorted traversal of all elements.

Problem: Merge K Sorted Lists

LeetCode 23 - Merge k Sorted Lists [hard]

Merge k sorted linked lists and return it as one sorted list. Analyze and describe its complexity.

Example:

Input:
[
  1->4->5,
  1->3->4,
  2->6
]
Output: 1->1->2->3->4->4->5->6

K-way Merge Solution

We need to find the smallest element of all the input lists. To do so, we have to compare only the smallest element of all the lists first. Once we have the smallest element, we can put it in the merged list.

Following a similar pattern, we can then find the next smallest element of all the lists to add it to the merged list.

1. We can push the smallest (first) element of each sorted array in a Min Heap to get the overall minimum.
2. After this, we can take out the smallest (top) element from the heap and add it to the merged list.
3. After removing the smallest element from the heap, we can insert the next element of the same list into the heap.
4. We can repeat steps 2 and 3 to populate the merged list in sorted order.

class ListNode(object):
    def __init__(self, x):
        self.val = x
        self.next = None


class ListNodeExtension(ListNode):
    def __lt__(self, other):
        return self.val < other.val


class Solution:
    def mergeKLists(self, lists: List[ListNode]) -> ListNode:
        ListNode.__lt__ = ListNodeExtension.__lt__
        min_heap = []
        for root in lists:
            if root is not None:
                heappush(min_heap, root)

        head = tail = ListNode(0)
        while min_heap:
            tail.next = heappop(min_heap)
            tail = tail.next
            if tail.next:
                heappush(min_heap, tail.next)

        return head.next

Time Complexity: O(N log K) where N is the total number of elements in all the K input arrays.

Space Complexity: O(K)

How to identify?

If the problem giving K sorted arrays and asks us to perform a sorted traversal of all the elements of all arrays, we need to think about K-way Merge pattern.

While solving the problems, we are going to use Heap data structure to keep track of all elements in K arrays.

Similar LeetCode Problems

• LeetCode 4 - Median of Two Sorted Arrays [hard]