DSA Factory
Free Heaps lessonsHeaps · Stage 2 · Merging streams · Step 3

The pairs with the smallest sums

List the k pairs, one value from each of two sorted arrays, with the smallest sums, by treating each value of the first array as the start of a sorted list of sums. About 18 minutes.

Each value starts a sorted row

You have two sorted arrays a and b. For one value a at position i, the sums with b are a plus b at position 0, a plus b at position 1, and so on: a sorted list. So the table of all sums is made of sorted rows, one per value of a.

The smallest sum overall is the front of one of these rows, so start the heap with the first pair of each row.

Row i: sums of a at i with every value of b.
Start: the first pair of each row.
In code
first = a[0] + b[0]   # the smallest sum
Sums of a = 2, 3, 9 with b = 1, 5, 8. Rows are values of a.
158
2
3
7
10
3
4
8
11
9
10
14
17

The heap starts with the first pair of each row: sums 3, 4 and 10. Pop 3: the pair (2, 1). Push the next in row 0: sum 7.

Move 1 of 4

Take a pair, offer the next

Pop the smallest entry. It is the next pair of the answer. Then offer the next pair in the same row, the one with the next value of b, if there is one.

Stop after k pairs or when the heap is empty. Because entries are compared by sum, then by position in a, then by position in b, ties are settled the same way every time. Only the first k rows can supply the answer, so start with at most k entries.

Pop: add that pair to the answer.
Push: the same row with j plus one.
In code
s, i, j = heapq.heappop(heap)
# the smallest sum
Quick check

After popping the pair at row i and column j, which pair is pushed?

  1. AThe pair in the same row, with the next value of b
  2. BThe pair in the next row, in the same column
  3. CAll pairs in the row
Show the answer

The pair in the same row, with the next value of b. It is the next sum in that sorted row.

Your problem

Find K pairs with the smallest sums

You are given two integer arrays a and b, both sorted in non-decreasing order, and an integer k. A pair takes one value from a and one from b. Return the k pairs with the smallest sums, listed in order of increasing sum; pairs with equal sums are listed by their position in a, then by their position in b. If there are fewer than k pairs, return all of them.

Example
a = [1, 7, 11], b = [2, 4, 6], k = 3 → [[1, 2], [1, 4], [1, 6]]

1 ≤ length of a, b ≤ 300 · -10,000 ≤ values ≤ 10,000 · 1 ≤ k ≤ 1,000

Solve it in your browserFree account, no card. Python, C++, Java or JavaScript, with hints if you get stuck, and your progress is saved.
Sign up free to solve