14. K-way Merge#
Use a min-heap (or divide-and-conquer merge) when you must combine k sorted streams — linked lists, matrix rows, or sorted arrays — and always pick the global minimum next element.
Template: push (value, stream_id, index) into a heap; pop smallest, append to result, push next element from that stream.
Problems at a glance#
| LC | Problem | |
|---|---|---|
| 23 | Merge k Sorted Lists | ↗ |
| 378 | Kth Smallest Element in a Sorted Matrix | ↗ |
| 632 | Smallest Range Covering Elements from K Lists | ↗ |
| 373 | Find K Pairs with Smallest Sums | ↗ |
| 786 | K-th Smallest Prime Fraction | ↗ |
| 313 | Super Ugly Number | ↗ |
Problems#
Problem 1: Merge k Sorted Lists (Leetcode:23)#
Problem Statement
You are given an array of k linked-lists lists, each linked-list is sorted in ascending order.
Merge all the linked-lists into one sorted linked-list and return it.
Example 1:
Input: lists = [[1,4,5],[1,3,4],[2,6]] Output: [1,1,2,3,4,4,5,6]
Constraints:
0 <= k <= 10^4, total nodes ≤ 10^4
Code and Explanation
import heapq
class Solution:
def mergeKLists(self, lists):
heap = []
for i, node in enumerate(lists):
if node:
heapq.heappush(heap, (node.val, i, node))
dummy = cur = ListNode(0)
while heap:
_, i, node = heapq.heappop(heap)
cur.next = node
cur = cur.next
if node.next:
heapq.heappush(heap, (node.next.val, i, node.next))
return dummy.next
- Seed the heap with the head of each non-empty list (store list index
ito break ties on equal values). - Repeatedly pop the smallest node, attach it to the result, and push its successor.
- Each node is pushed and popped once.
Time: O(N log k) where N = total nodes · Space: O(k) heap
class Solution:
def mergeKLists(self, lists):
if not lists:
return None
def merge_two(a, b):
dummy = cur = ListNode(0)
while a and b:
if a.val <= b.val:
cur.next, a = a, a.next
else:
cur.next, b = b, b.next
cur = cur.next
cur.next = a or b
return dummy.next
while len(lists) > 1:
merged = []
for i in range(0, len(lists), 2):
l1 = lists[i]
l2 = lists[i + 1] if i + 1 < len(lists) else None
merged.append(merge_two(l1, l2))
lists = merged
return lists[0]
Pairwise merge lists like merge sort. log k rounds, each touching all N nodes → O(N log k). No heap; good when heap API is awkward.
Time: O(N log k) · Space: O(1) extra (iterative) or O(log k) recursion
Problem 2: Kth Smallest Element in a Sorted Matrix (Leetcode:378)#
Problem Statement
Given an n x n matrix where each row and column is sorted in ascending order, return the kth smallest element in sorted order.
Example 1:
Input: matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8 Output: 13
Code and Explanation
import heapq
class Solution:
def kthSmallest(self, matrix: list[list[int]], k: int) -> int:
n = len(matrix)
heap = [(matrix[r][0], r, 0) for r in range(n)]
heapq.heapify(heap)
for _ in range(k - 1):
val, r, c = heapq.heappop(heap)
if c + 1 < n:
heapq.heappush(heap, (matrix[r][c + 1], r, c + 1))
return heapq.heappop(heap)[0]
Each row is a sorted stream. Start with column 0 of every row in the heap. Each pop advances that row by one column. After k-1 expansions, the next pop is the kth smallest.
Time: O(k log n) · Space: O(n)
class Solution:
def kthSmallest(self, matrix: list[list[int]], k: int) -> int:
n = len(matrix)
def count_le(x):
r, c, cnt = n - 1, 0, 0
while r >= 0 and c < n:
if matrix[r][c] <= x:
cnt += r + 1
c += 1
else:
r -= 1
return cnt
lo, hi = matrix[0][0], matrix[-1][-1]
while lo < hi:
mid = (lo + hi) // 2
if count_le(mid) >= k:
hi = mid
else:
lo = mid + 1
return lo
Not k-way merge, but optimal for large k: binary search the answer value. Count elements ≤ mid by walking from bottom-left (staircase). First mid with count ≥ k is the answer.
Time: O(n log(max−min)) · Space: O(1)
Problem 3: Smallest Range Covering Elements from K Lists (Leetcode:632)#
Problem Statement
You have k lists of sorted integers. Find the smallest range [a, b] that includes at least one number from each list.
Example 1:
Input: nums = [[4,10,15,24,26],[0,9,12,20],[5,18,22,30]] Output: [20,24]
Code and Explanation
import heapq
class Solution:
def smallestRange(self, nums: list[list[int]]) -> list[int]:
k = len(nums)
heap = [(row[0], i, 0) for i, row in enumerate(nums)]
heapq.heapify(heap)
cur_max = max(row[0] for row in nums)
best = (-10**9, 10**9)
while len(heap) == k:
cur_min, r, c = heapq.heappop(heap)
if cur_max - cur_min < best[1] - best[0]:
best = (cur_min, cur_max)
if c + 1 == len(nums[r]):
break
nxt = nums[r][c + 1]
cur_max = max(cur_max, nxt)
heapq.heappush(heap, (nxt, r, c + 1))
return [best[0], best[1]]
Track current window [min in heap, max seen]. Each step advances the list that contributed the minimum. Stop when any list is exhausted — no valid range can include all lists after that.
Time: O(N log k) · Space: O(k)
Problem 4: Find K Pairs with Smallest Sums (Leetcode:373)#
Problem Statement
Given sorted nums1, nums2, return k pairs (u, v) with smallest sums.
Example 1:
Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3 Output: [[1,2],[1,4],[1,6]]
Code and Explanation
import heapq
class Solution:
def kSmallestPairs(self, nums1: list[int], nums2: list[int], k: int) -> list[list[int]]:
if not nums1 or not nums2:
return []
heap = [(nums1[i] + nums2[0], i, 0) for i in range(min(k, len(nums1)))]
heapq.heapify(heap)
res = []
while heap and len(res) < k:
_, i, j = heapq.heappop(heap)
res.append([nums1[i], nums2[j]])
if j + 1 < len(nums2):
heapq.heappush(heap, (nums1[i] + nums2[j + 1], i, j + 1))
return res
Fix nums1[i] and treat pairs (nums1[i], nums2[j]) as a sorted stream over j. Seed heap with (nums1[i]+nums2[0], i, 0) for first k rows — only these can beat larger i initially. Each pop advances j for that row.
Time: O(k log k) · Space: O(k)
class Solution:
def kSmallestPairs(self, nums1, nums2, k):
lo, hi = nums1[0] + nums2[0], nums1[-1] + nums2[-1]
def count_le(s):
cnt, j = 0, len(nums2) - 1
for x in nums1:
while j >= 0 and x + nums2[j] > s:
j -= 1
cnt += j + 1
return cnt
while lo < hi:
mid = (lo + hi) // 2
if count_le(mid) >= k:
hi = mid
else:
lo = mid + 1
# collect pairs with sum <= lo, trim to k
...
Binary search the kth smallest sum, then collect pairs. Useful when k approaches m×n; heap is simpler for typical interview k.
Time: O((m+n) log S) · Space: O(k)
Problem 5: K-th Smallest Prime Fraction (Leetcode:786)#
Problem Statement
Sorted array of unique primes (with leading 1). Return the kth smallest fraction arr[i]/arr[j] where i < j.
Example 1:
Input: arr = [1,2,3,5], k = 3 Output: [2,5]
Code and Explanation
import heapq
class Solution:
def kthSmallestPrimeFraction(self, arr: list[int], k: int) -> list[int]:
n = len(arr)
heap = [(arr[0] / arr[j], 0, j) for j in range(1, n)]
heapq.heapify(heap)
for _ in range(k - 1):
_, i, j = heapq.heappop(heap)
if i + 1 < j:
heapq.heappush(heap, (arr[i + 1] / arr[j], i + 1, j))
_, i, j = heapq.heappop(heap)
return [arr[i], arr[j]]
For fixed denominator index j, fractions arr[i]/arr[j] increase as i increases. Seed heap with (arr[0]/arr[j], 0, j) for each j ≥ 1. Each pop tries the next numerator for the same denominator.
Time: O(k log n) · Space: O(n)
class Solution:
def kthSmallestPrimeFraction(self, arr, k):
def count_le(x):
cnt, j = 0, len(arr) - 1
for i in range(len(arr)):
while j > i and arr[i] / arr[j] > x:
j -= 1
cnt += max(0, j - i)
return cnt
# binary search float x, track best pair ...
Count pairs with fraction ≤ x using two pointers per row. Binary search x until count ≥ k. Avoids heap; more code for exact pair recovery.
Time: O(n log(max/min) · log precision) · Space: O(1)
Problem 6: Super Ugly Number (Leetcode:313)#
Problem Statement
Return the nth super ugly number — positive integers whose prime factors all lie in primes.
Example 1:
Input: n = 12, primes = [2,7,13,19] Output: 32
Code and Explanation
import heapq
class Solution:
def nthSuperUglyNumber(self, n: int, primes: list[int]) -> int:
seen = {1}
heap = [1]
for _ in range(n):
ugly = heapq.heappop(heap)
for p in primes:
nxt = ugly * p
if nxt not in seen:
seen.add(nxt)
heapq.heappush(heap, nxt)
return ugly
Each ugly number generates ugly * p for every prime p. Min-heap always yields the next smallest unseen ugly. Dedup set prevents duplicates (e.g. 2×3 vs 3×2).
Time: O(n log n · k) roughly · Space: O(n)
class Solution:
def nthSuperUglyNumber(self, n: int, primes: list[int]) -> int:
ugly = [1]
ptr = [0] * len(primes)
for _ in range(1, n):
nxt = min(ugly[ptr[i]] * primes[i] for i in range(len(primes)))
ugly.append(nxt)
for i in range(len(primes)):
if ugly[ptr[i]] * primes[i] == nxt:
ptr[i] += 1
return ugly[-1]
Same idea as LC 264 (Ugly Number II): one pointer per prime stream into the ugly array. No heap; each pointer only moves forward. Standard follow-up when n is large.
Time: O(n · k) · Space: O(n)