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
- if row is non empty, then
- 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
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]