Find K Pairs With Smallest Sums: Frontier Expansion

"Find K Pairs with Smallest Sums" shows up again and again in Uber, Amazon, Google phone screens. It is a Medium problem on paper, but the real test is whether you recognize the Top-K and Heaps pattern quickly and code it cleanly. Part 9 of 11 in the Top-K and Heaps arc.

The Problem

Given two sorted arrays and an integer k, return the k pairs (one element from each) with the smallest sums. The pair space is a virtual matrix of n·m sums, sorted along both axes — materializing it is the trap; the heap explores just the frontier.

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

Recognizing the Top-K and Heaps Pattern

When a problem asks for the k largest, smallest, closest, or most frequent — or for repeated access to an extreme while data changes — a heap gives O(log n) insertion and O(1) access to the extreme. The signature trick: keep a bounded heap of size k with the opposite polarity (a min-heap to track the k largest), evicting the root on overflow. Python's heapq is a min-heap; negate values for max behavior.

K smallest from a structure sorted along multiple dimensions: a min-heap holds the frontier — candidates whose predecessors are already emitted — and each pop admits its successors. Seeding pairs (i, 0) for the first min(n, k) rows and advancing j on pop keeps every candidate reachable exactly once.

The Approach

Push (nums1[i] + nums2[0], i, 0) for the first min(n, k) values of i. Pop the smallest sum, emit the pair, and push its row successor (i, j + 1) if it exists. K pops with O(log k) heap operations each.

Why no duplicate visits: every pair (i, j) has exactly one admission path — from (i, j − 1) — except the seeded column. Being able to answer that is the difference between using the technique and understanding it; the same walk solves Kth Smallest in a Sorted Matrix verbatim.

Python Solution

import heapq


def k_smallest_pairs(
    nums1: list[int], nums2: list[int], k: int
) -> list[list[int]]:
    """K pairs with smallest sums from two sorted arrays."""
    if not nums1 or not nums2 or k <= 0:
        return []
    heap: list[tuple[int, int, int]] = []
    for i in range(min(len(nums1), k)):
        heapq.heappush(heap, (nums1[i] + nums2[0], i, 0))

    result: list[list[int]] = []
    while heap and len(result) < k:
        _, i, j = heapq.heappop(heap)
        result.append([nums1[i], nums2[j]])
        if j + 1 < len(nums2):
            heapq.heappush(heap, (nums1[i] + nums2[j + 1], i, j + 1))
    return result

Complexity

• Time: O(k log k) — k pops and at most k + min(n, k) pushes
• Space: O(k) — frontier heap and output

Interview Tips and Follow-Ups

• Seeding all n rows when k is small wastes O(n) — the min(n, k) bound matters and interviewers check for it.
• Kth Smallest in a Sorted Matrix is this exact frontier walk over rows — one problem, two entries on the question lists.
• The alternative visited-set formulation (push both (i+1, j) and (i, j+1), dedupe) also works — compare briefly and justify your pick.

If this clicked, continue the Top-K and Heaps arc in the Technical Interview category. One hundred questions, nine patterns, all in Python.