Medium K-way Merge Heaps
In plain English: Combine k already-sorted lists into one single sorted list efficiently.
A min-heap of size k always holds the smallest unprocessed element from each array — pop the min and push the next from that array.
Prompt
Given k sorted arrays, merge them into one sorted array.
Try to write it from scratch before scrolling down.
Solution
import heapq
def merge_k_sorted(arrays):
heap = []
for i, arr in enumerate(arrays):
if arr:
heapq.heappush(heap, (arr[0], i, 0)) # (val, array_idx, elem_idx)
result = []
while heap:
val, i, j = heapq.heappop(heap)
result.append(val)
if j + 1 < len(arrays[i]):
heapq.heappush(heap, (arrays[i][j + 1], i, j + 1))
return result
# Test: merge_k_sorted([[1,4,5],[1,3,4],[2,6]]) == [1,1,2,3,4,4,5,6]