K-Way Merge

Posted on

min_heap 的应用，特别之处在于借助 index 对元素进行 track。

This pattern helps us solve problems that involve a list of sorted arrays.

Whenever we are given “K” sorted arrays, we can use a min heap to efficiently perform a sorted traversal of all the elements of all arrays. We can push the smallest (first) element of each sorted array in a min heap to get the overall minimum. While inserting elements to the min heap we keep track of which array the element came from. We can, then, 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.

Snippets

1
2
3
4
5
6
7
8
9
10
11
min_heap = []
for row in matrix:
    heappush(min_heap, (row[0], 0, row))  # keep track by index

n = len(matrix)
for _ in range(k - 1):
    _, idx, row = heappop(min_heap)
    if idx + 1 < n:
        heappush(min_heap, (row[idx+1], idx + 1, row))

return min_heap[0][0]

LeetCode

23. Merge k Sorted Lists
4. Median of Two Sorted Arrays
378. Kth Smallest Element in a Sorted Matrix
632. Smallest Range Covering Elements from K Lists
373. Find K Pairs with Smallest Sums