Find K Pairs with Smallest Sums (LeetCode 373)

The problem

You get two arrays nums1 and nums2, both sorted in non-decreasing order, and an integer k. A pair takes one number from each array. Return the k pairs with the smallest sums.

Input:  nums1 = [1, 7, 11], nums2 = [2, 4, 6], k = 3
Output: [[1, 2], [1, 4], [1, 6]]

The sums are 3, 5 and 7. Every pair that uses 7 or 11 is at least 9, so all three answers use 1.

How to think about it

Start with brute force. Build all len(nums1) × len(nums2) pairs, sort by sum, and take the first k. That’s O(mn log mn), and m and n can each be 10⁵, so that’s 10¹⁰ pairs. You can’t even list them.

Picture the pairs as a grid. Row i holds nums1[i] + nums2[0], nums1[i] + nums2[1], and so on. Because nums2 is sorted, every row is already a sorted list. Finding the k smallest across all rows is exactly merging k sorted lists, and you only need the first k merged items.

The key insight: you only need the current front of each row. Put the first item of each row, (nums1[i] + nums2[0], i, 0), in a min-heap. Pop the smallest; that’s the next answer. Then push its row’s next item, (i, j + 1). The heap always holds the smallest unused item of each active row, so the next answer is always at its root. And since only k answers are needed, only the first k rows can ever matter, so seed just min(k, m) of them.

Which pattern is this? Heaps, as a k-way merge. The tell is “k smallest” combined with inputs that are already sorted, so candidates can be produced lazily instead of all at once.

Approach

1. For each i in the first min(k, len(nums1)) rows, push (nums1[i] + nums2[0], i, 0).
2. While the heap isn’t empty and you have fewer than k pairs:
   • pop (s, i, j) and record [nums1[i], nums2[j]];
   • if j + 1 is a valid index, push (nums1[i] + nums2[j + 1], i, j + 1).
3. Return the recorded pairs.

Solution

import heapq
def k_smallest_pairs(nums1, nums2, k):
    h = []; res = []
    for i in range(min(k, len(nums1))):
        heapq.heappush(h, (nums1[i] + nums2[0], i, 0))      # front of each row
    while h and len(res) < k:
        s, i, j = heapq.heappop(h); res.append([nums1[i], nums2[j]])   # next smallest overall
        if j + 1 < len(nums2): heapq.heappush(h, (nums1[i] + nums2[j+1], i, j+1))   # advance along that row
    return res
# k_smallest_pairs([1,7,11], [2,4,6], 3) -> [[1,2],[1,4],[1,6]]

Dry run

On the example. The heap is seeded with (3, 0, 0), (9, 1, 0), (13, 2, 0):

Three pairs are recorded, so the loop stops: [[1, 2], [1, 4], [1, 6]].

Complexity

• Time: O(k log k). Seeding pushes at most k items. Then each of the k pops may push one item, and the heap never holds more than k items.
• Space: O(k) for the heap and the result.

Edge cases and pitfalls

• k larger than the number of pairs: [1, 2], [3], k = 3 returns just the two pairs that exist. The while h check handles it.
• Duplicates are fine: [1, 1, 2], [1, 2, 3], k = 2 gives [[1, 1], [1, 1]], from different indices.
• Pushing both (i+1, j) and (i, j+1) from each pop is another valid approach, but then you need a visited set to avoid pushing the same cell twice.
• Seeding every row when nums1 is huge wastes time; cap it at k.
• Empty input arrays would crash on nums2[0]; LeetCode guarantees both are non-empty, but guard it if yours aren’t.

Related problems

• Merge k Sorted Lists (LeetCode 23): the k-way merge this solution is built on.
• K Closest Points to Origin (LeetCode 973): k smallest when every candidate is already known.
• Kth Largest Element in an Array (LeetCode 215): the basic top-K heap.