Generated by Codex with GPT 5.6 Sol XHigh

Quick facts

Problem gist

Two sorted integer arrays are given. A pair is formed by choosing one value from the first array and one from the second, and its cost is their sum. The task is to return the k pairs with the smallest costs, or every possible pair when fewer than k pairs exist.

Generating all m * n pairs and sorting them is wasteful when k is small. The sorted inputs provide much more structure: once one value is fixed, pairing it with the other array produces an already sorted sequence of sums. The problem is therefore a merge of sorted sequences rather than a full Cartesian-product sort.

Deriving the optimal solution

Imagine a matrix in which each cell contains one pair sum. A row fixes one value from one array and combines it with every value in the other array. Because both arrays are sorted, every row is sorted from left to right. The implementation uses the shorter array for the rows, which keeps the heap as small as possible; if the arrays are transposed, it restores each result to the required [nums1 value, nums2 value] order.

The first cell of every active row is placed in a min-heap. The heap then contains the smallest unseen pair from each row:

  1. Remove the heap’s smallest pair and add it to the answer.
  2. Advance one position in that same row.
  3. Insert the row’s new frontier into the heap.

This is the same idea used to merge multiple sorted lists. If the heap minimum is at row i, column j, every earlier cell in that row has already been returned, every later cell in that row is at least as large, and the heap still represents the smallest unseen cell of every other active row. The removed cell is therefore the globally smallest unseen pair.

Only the first k rows can matter. The first cell of row k has at least k row-starting cells before it that are no larger, so a later row cannot be necessary to produce the first k results. Ties do not change this conclusion: choosing any k pairs at the cutoff sum is valid. Thus the number of heap entries is

h = min(k, len(nums1), len(nums2)).

Building the heap takes O(h) time. If r = min(k, m * n) pairs are returned, the repeated heap operations take O(r log h) time. Auxiliary space is O(h), while the required output occupies O(r) space. This avoids doing work proportional to all m * n combinations when only a small prefix is requested.

Python solution

from __future__ import annotations

from heapq import heapify, heappop, heappush
from typing import Sequence, TypeAlias


HeapEntry: TypeAlias = tuple[int, int, int]


class KSmallestPairFinder:
    """Find the lowest-sum pairs from two sorted integer sequences."""

    @staticmethod
    def find(
        first_values: Sequence[int],
        second_values: Sequence[int],
        pair_limit: int,
    ) -> list[list[int]]:
        """Return up to pair_limit pairs in nondecreasing sum order.

        Both input sequences must be sorted in nondecreasing order. The inputs
        are read only and are never copied or modified.
        """
        if pair_limit <= 0 or not first_values or not second_values:
            return []

        # Treat the shorter input as the set of rows to minimize heap space.
        # When the arrays are transposed, restore the original pair order when
        # appending to the result.
        if len(first_values) <= len(second_values):
            row_values = first_values
            column_values = second_values
            arrays_transposed = False
        else:
            row_values = second_values
            column_values = first_values
            arrays_transposed = True

        active_row_count = min(pair_limit, len(row_values))
        candidate_heap: list[HeapEntry] = [
            (row_values[row_index] + column_values[0], row_index, 0)
            for row_index in range(active_row_count)
        ]
        heapify(candidate_heap)

        smallest_pairs: list[list[int]] = []

        while candidate_heap and len(smallest_pairs) < pair_limit:
            _, row_index, column_index = heappop(candidate_heap)

            if arrays_transposed:
                pair = [column_values[column_index], row_values[row_index]]
            else:
                pair = [row_values[row_index], column_values[column_index]]
            smallest_pairs.append(pair)

            # Only the next pair in this row can become its new minimum unseen
            # candidate; all later pairs are at least as large.
            next_column_index = column_index + 1
            if next_column_index < len(column_values):
                next_sum = (
                    row_values[row_index]
                    + column_values[next_column_index]
                )
                heappush(
                    candidate_heap,
                    (next_sum, row_index, next_column_index),
                )

        return smallest_pairs


class Solution:
    """LeetCode-compatible entry point."""

    def kSmallestPairs(
        self,
        nums1: list[int],
        nums2: list[int],
        k: int,
    ) -> list[list[int]]:
        return KSmallestPairFinder.find(nums1, nums2, k)

Interview follow-ups

Why is it safe to initialize only the first k rows?

The first cell in each row is that row’s minimum. Because the row values are sorted, the minimum of row k is no smaller than the minima of rows 0 through k - 1. Those earlier rows alone already provide k pairs whose sums are no larger, so no pair from row k or beyond is required among the first k results.

Equal sums do not break the proof. If an omitted row begins at the same cutoff sum as several included rows, the problem permits any k pairs with the smallest sums. Limiting the initial heap therefore reduces its size without changing correctness, giving O(min(k, m, n)) auxiliary space after orienting the shorter array as the rows.

What changes if the input arrays are not sorted?

Sort copies of both arrays first, then run the same heap merge. Sorting costs O(m log m + n log n) time and O(m + n) space for copied inputs in Python; the pair-selection phase keeps its existing bound. If mutation is acceptable, in-place sorting can reduce the extra array storage.

If the result must include original indices, sort (value, original_index) records instead of bare values. The heap should compare pair sums built from the values while the returned records retain their original positions. This preserves the algorithm but adds index metadata to every stored element and output pair.

How would the solution return only the k-th smallest sum?

When only the cutoff value is needed, binary-search a candidate sum rather than materializing k pairs. For a proposed sum S, count pairs with total at most S using two pointers: start at the largest value of one array and move monotonically through the other. The count takes O(m + n) time because neither pointer reverses direction.

Binary search over the range from the smallest possible sum to the largest possible sum. The first value with at least k qualifying pairs is the k-th smallest sum, including duplicate sums. For integer values, this takes O((m + n) log R) time, where R is the numeric sum range, and O(1) auxiliary space. It is attractive for very large k, but it does not identify the actual pairs without an additional enumeration step.

What if the output must contain unique value pairs?

First compress each sorted array to its distinct values, then apply the same row-frontier heap algorithm to the compressed arrays. Every matrix cell now represents a distinct ordered value pair, so no output-level hash set is needed. Compression takes O(m + n) time, after which the heap size depends on the number of distinct values rather than the original lengths.

This interpretation must be clarified with the interviewer. The original problem treats different index combinations as possible pairs even when their values are equal, whereas “unique pairs” usually means uniqueness by the two returned values. If uniqueness instead means unique source indices, the original algorithm already distinguishes positions and should not compress anything.

Can the pairs be produced as a stream instead of a list?

Yes. The heap loop can yield one pair after each pop rather than appending to an output list. The frontier invariant and O(log h) delay per emitted pair stay unchanged, while memory becomes O(h) because the caller consumes the results incrementally.

Streaming is useful when the caller may stop early or send results over a network. The tradeoff is interface complexity: errors or cancellation can occur partway through iteration, and a caller that ultimately retains every pair still uses O(k) memory outside the generator. LeetCode’s required return type is a list, so the production entry point above materializes the output.

How would the approach change for three sorted arrays?

The sums form a three-dimensional sorted grid. A best-first search can start at index triple (0, 0, 0), repeatedly remove the smallest sum, and add its three forward neighbors. A visited set is required because the same index triple can be reached along different paths. This produces k results in roughly O(k log k) time and O(k) auxiliary space.

Another approach first generates a bounded set of small pairs from two arrays and then merges those sums with the third array. That can reuse the two-array routine, but the intermediate bound needs a careful proof so a necessary combination is not discarded. The grid search is usually easier to justify in an interview, at the cost of a larger visited-state structure than the two-array row merge.

What if k is close to the total number of pairs?

Any algorithm that returns the pairs explicitly already needs Theta(k) time and output space, so there is no sublinear shortcut when k approaches m * n. The heap merge remains correct and avoids sorting the full product, with O(k log min(m, n)) time after the heap is initialized.

If output order is irrelevant and almost every pair is required, direct nested-loop enumeration can be faster in practice because it removes heap overhead, though it does not naturally stop at the exact smallest subset when k < m * n. The interviewer should expect the choice to depend on whether sorted output is required and how close k is to the full Cartesian-product size.