K-Way Merge · bigonotes

When you have K sorted sources (lists, matrix rows, streams), the next smallest item overall must be the head of one of them. A K-way merge keeps only those K heads in a min-heap: pop the smallest, then push the next item from the same source. It shows up in merging, in "Kth smallest" over sorted rows, and in "K smallest pairs".

The idea#

Each source is sorted, so its head is the smallest item it has left. The global minimum is therefore the smallest of the K heads, which is exactly what the heap's top gives you.

Push the first item of every source, tagged with where it came from.
Pop the smallest. It is the next item of the merged output.
Push the next item from that source, if there is one.

The heap never holds more than K items, so each step is $O(\log K)$ . Merging N items in total costs $O(N \log K)$ . Pushing every item into one heap also works, but it costs $O(N \log N)$ time and $O(N)$ space, and it ignores the fact that the sources are already sorted.

To find the Kth smallest item, stop after K pops instead of merging everything.

How to spot it#

These cues suggest a K-way merge.

"Merge K sorted lists / arrays".
A matrix whose rows (or columns) are sorted, and you need the Kth smallest or the median.
Pairs or sums from two sorted arrays: each nums1[i] defines a sorted "row" over nums2.
Numbers generated by multiplying earlier ones (Ugly Number II): several sorted streams being merged.

Template#

Here are the three steps as code, for any number of sorted sources. The source index sits in the tuple, so ties never fall through to comparing the items themselves.

heap = (first item, source i, position 0) for each non-empty source
heapify heap                              # one head per source: O(K)
while heap is not empty (and fewer than limit items produced):
    pop (value, i, j) from heap
    output value
    if source i has an item at j + 1:     # advance the same source only
        push (source i's item at j + 1, i, j + 1)

As in "The idea", stop after K pops when you only need the Kth smallest. One more detail: some payloads, such as linked list nodes, can't be compared, so a tie between two (val, node) tuples raises an error. Either keep the index in the tuple as above, or wrap the node in a small class whose __lt__ compares the node values.

Pitfalls#

Most bugs come from losing track of which source an item belongs to.

Pushing every item. It's correct, but it throws away the sortedness and costs $O(N \log N)$ .
Comparing payloads on ties. Put a source index (or a counter) before the payload, or define __lt__.
Advancing the wrong source. After popping from source i, push the next item from source i only.
Over-seeding. For "Kth" questions, seed only min(K, number of sources) heads. Sources beyond the Kth can't hold any of the first K items.

Problems in this pattern#

These problems use this pattern. Solve them before you open the hints below. Start with the first one.

Hints#

Open a hint only after you have tried the problem for a while.

Hint: Merge k Sorted Lists

The heads are list nodes, which can't be compared. Push (val, index, node), or wrap each node in a class with __lt__. After a pop, push that node's next if it exists.

Hint: Kth Smallest Element in a Sorted Matrix

Treat each row as a sorted list and pop K times. Seed at most min(K, n) rows. That is about $O(K \log n)$ , against $O(n^2 \log K)$ for scanning the whole matrix with a size-K heap. Binary search on the value range is the other approach; see Searching 2D Matrices.

Hint: Find K Pairs with Smallest Sums

Each nums1[i] defines a sorted row over nums2. Seed (nums1[i] + nums2[0], i, 0) for i < min(K, len(nums1)). Each pop records a pair and pushes (i, j + 1), so you touch only about K pairs instead of all $mn$ .

Hint: Median of a Row Wise Sorted Matrix

For an odd-sized m × n matrix, the median is the (m * n) // 2 + 1th smallest value. Pop that many times. Binary search on the value is faster for large matrices.

Hint: Ugly Number II

Each ugly number comes from multiplying an earlier one by 2, 3 or 5. Treat those products as sorted streams being merged, and skip duplicates.

Key takeaways#

The next smallest item overall is the smallest of the K heads.
Keep one head per source: $O(N \log K)$ to merge everything, or stop after K pops for the Kth smallest.
After a pop, push the successor from the same source.
Add an index to the tuple, or define __lt__, so the heap never compares payloads.
A size-K heap (Top K) and a K-way merge are different patterns: one caps the heap at K items, the other keeps one item per source.

Core idea#

With K sorted sources, the next smallest item overall is the head of one of them. Keep only the K heads in a min-heap. Pop the smallest, then push the next item from the same source.

Key points#

Seed the heap with the first item of each source, tagged with its source.
The heap holds at most K items, so each step is $O(\log K)$ . Merging N items costs $O(N \log K)$ .
Pushing every item into one heap costs $O(N \log N)$ time and $O(N)$ space.
Kth smallest: stop after K pops.
Sorted matrix: each row is a source. Seed at most min(K, n) rows. Cost is about $O(K \log n)$ .
K smallest pairs: seed (nums1[i] + nums2[0], i, 0). On a pop, push (i, j + 1).
Median of a row-sorted matrix: pop (m * n) // 2 + 1 times.
Spot it: merge K sorted lists, Kth smallest in sorted rows, pairs from two sorted arrays, Ugly Number II.

Template#

heap = (first item, source i, position 0) for each non-empty source
heapify heap                              # one head per source
while heap is not empty:
    pop (value, i, j) from heap
    output value
    if source i has an item at j + 1:     # same source only
        push (source i's item at j + 1, i, j + 1)

Pitfalls#

Pushing every item. It ignores the sortedness.
Comparing payloads on ties. A (val, node) tie fails for ListNode. Add a source index or define __lt__.
Advancing the wrong source.
Over-seeding for "Kth" questions.

Say it in an interview#

"The next smallest item is the smallest of the K heads, so the heap holds one per source."
"Top K caps the heap at K items. K-way merge keeps one item per source."