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 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
Merge k sorted linked lists and return it as one sorted list. Analyze and describe its complexity.
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.
- We can push the smallest (first) element of each sorted array in a Min Heap to get the overall minimum.
- After this, we can take out the smallest (top) element from the heap and add it to the merged list.
- After removing the smallest element from the heap, we can insert the next element of the same list into the heap.
- 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.