Program to merge K-sorted lists in Python

Suppose we have some lists, these lists are sorted. We have to merge these lists into one list. To solve this, we will use the heap data structure. So if the lists are [1,4,5], [1,3,4], [2,6], then the final list will be [1,1,2,3,4,4,5,6].

To solve this, we will follow these steps −

• n := size of lists
• heap := a new list
• for each index i and row of lists[i], do
  • if row is non empty, then
    • insert (row[0], i, 0) into heap
• res := a new list
• while heap is not empty, do
  • num, row, col := top element of heap
  • insert num at the end of res
  • if col < size of lists[row] - 1, then
    • insert lists[row, col + 1], row, col + 1 into heap
• return res

Let us see the following implementation to get better understanding −

Example

 Live Demo

import heapq
class Solution:
   def solve(self, lists):
      n = len(lists)
      heap = []
      for i, row in enumerate(lists):
         if row:
            heapq.heappush(heap, (row[0], i, 0))
         res = []
         while heap:
            num, row, col = heapq.heappop(heap)
            res.append(num)
         if col < len(lists[row]) - 1:
            heapq.heappush(heap, (lists[row][col + 1], row, col + 1))
      return res
ob = Solution()
lists = [[],[],[11, 13],[],[4, 4, 14],[4],[11],[1, 8]]
print(ob.solve(lists))

Input

[[],[],[11, 13],[],[4, 4, 14],[4],[11],[1, 8]]

Output

[1, 4, 4, 4, 8, 11, 11, 13, 14]