K-Way Merge
Two Heaps & K-Way Merge: lesson 2 of 4
One heap of k heads turns k sorted lists into one.
Lesson 2 of 4 · 6 min
K-Way Merge
Step 1 of 17
in heap0
out
Three lanes, each already sorted. Only the value at the front of a lane can be the next smallest overall.
The Idea
Only the front value of a sorted list can be the next smallest overall. So hold just those k heads in a min-heap.
Pop the winner, append it to the output, and push the value that stepped up behind it. The heap never grows past k, no matter how long the lists are.
Real-World Example
A postal sorting hall with k conveyor belts, each already in date order. A clerk only ever looks at the k parcels at the belt fronts, takes the oldest, and the belt behind it rolls one step forward.
The Code
import heapq
def merge_k(lists):
heap = [(rows[0], i, 0) for i, rows in enumerate(lists) if rows]
heapq.heapify(heap) # k heads, nothing more
out = []
while heap:
value, i, j = heapq.heappop(heap)
out.append(value)
if j + 1 < len(lists[i]): # advance only that list
heapq.heappush(heap, (lists[i][j + 1], i, j + 1))
return out
print(merge_k([[1, 6, 9], [2, 4], [3, 8]])) # [1, 2, 3, 4, 6, 8, 9]
print(merge_k([[], [5]])) # [5]Your turn
What does this print?
import heapq
heap = [(1, 0, 0), (2, 1, 0), (3, 2, 0)] # the three heads
lists = [[1, 6], [2, 4], [3, 8]]
value, i, j = heapq.heappop(heap)
heapq.heappush(heap, (lists[i][j + 1], i, j + 1))
print(value, heap[0][0])Mini quiz
1 / 3