~/problems / Heaps / Heaps and priority queues

Find K Pairs with Smallest Sums

On a phone? Coding is easier on a laptop: email this problem to yourself . Meanwhile: quiz this topic or fight a boss.

medium ~25 min

nums1 and nums2 are lists of integers, each sorted from smallest to largest. A pair [u, v] takes one entry u from nums1 and one entry v from nums2, so there are len(nums1) * len(nums2) pairs. Pairs are told apart by the positions they use: if a value appears twice in nums1, each copy makes its own pairs.

Write k_smallest_pairs(nums1: list[int], nums2: list[int], k: int) -> list[list[int]] that returns the k pairs with the smallest sums u + v. If there are fewer than k pairs, return all of them. The pairs may come in any order, and when several pairs tie for the last places, any of them will do.

Examples:

k_smallest_pairs([1, 4, 6], [2, 3, 9], 4)   # [[1, 2], [1, 3], [4, 2], [4, 3]]

The 9 sums are 3, 4, 10, 6, 7, 13, 8, 9, 15 (pairing 1, then 4, then 6 with each of 2, 3, 9). The four smallest are 3, 4, 6 and 7. Notice that [1, 9] (sum 10) loses to [4, 3] (sum 7), even though 1 is the smallest value in nums1.

k_smallest_pairs([0, 0], [5, 7], 3)         # [[0, 5], [0, 5], [0, 7]]

Each 0 makes its own [0, 5], so that pair appears twice. The third place is a tie between the two [0, 7] pairs, which look the same.

k_smallest_pairs([-2], [3, 4], 5)           # [[-2, 3], [-2, 4]]

There are only two pairs, so both come back.

Constraints: 1 <= len(nums1), len(nums2) <= 10^5, values -10^9..10^9, 1 <= k <= 10^4.

⭐ Bonus: a speed test with two lists of 100,000 values and k = 10,000 earns a star. That's 10 billion pairs, so listing them all (or even all pairs of the first k values of each list) is far too slow.

Show hint

Line the pairs up as a grid: row i pairs nums1[i] with nums2[0], nums2[1], ..., and each row's sums only grow. That's sorted lists again: keep the front of each row in a min-heap and take the smallest k times.

Topic: Heaps and priority queues. heapq: top-k, k-way merge, two heaps for a running median.

Read the visual guide
0:00
Ctrl ' run · Ctrl ↵ submit
esc