Find K Pairs With the Smallest Sums

Given two arrays nums1 and nums2 sorted in ascending order and an integer k, return the k pairs (u, v), with u from nums1 and v from nums2, that have the smallest sums.

LeetCode: 373. Find K Pairs with Smallest Sums

Example

Text

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

nums1 = [1, 1, 2], nums2 = [1, 2, 3], k = 2
Output: [(1, 1), (1, 1)]

Brute force: sort all pairs

Java

import java.util.ArrayList;
import java.util.List;

public class Main {
    static List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<int[]> pairs = new ArrayList<>();
        for (int u : nums1) {
            for (int v : nums2) {
                pairs.add(new int[]{u, v});
            }
        }
        // List.sort is stable: ties keep their generation order
        pairs.sort((a, b) -> Long.compare((long) a[0] + a[1], (long) b[0] + b[1]));
        return new ArrayList<>(pairs.subList(0, Math.min(k, pairs.size())));
    }

    static String format(List<int[]> pairs) {
        StringBuilder sb = new StringBuilder("[");
        for (int i = 0; i < pairs.size(); i++) {
            if (i > 0) sb.append(", ");
            sb.append("(").append(pairs.get(i)[0]).append(", ").append(pairs.get(i)[1]).append(")");
        }
        return sb.append("]").toString();
    }

    public static void main(String[] args) {
        System.out.println(format(kSmallestPairs(new int[]{1, 7, 11}, new int[]{2, 4, 6}, 3)));
        System.out.println(format(kSmallestPairs(new int[]{1, 1, 2}, new int[]{1, 2, 3}, 2)));
    }
}

Output

[(1, 2), (1, 4), (1, 6)]
[(1, 1), (1, 1)]

Time: O(mn log(mn)). Space: O(mn). With arrays of 100,000 elements, that's 10¹⁰ pairs.

Optimal: a min-heap of candidates

Picture the sums as a grid: row i holds nums1[i] + nums2[j] for every j. Every row is sorted, so the smallest unused pair of a row is always the leftmost one not taken yet. That makes this a merge of sorted rows:

Push the first pair of each row, (i, 0), onto a min-heap. Only the first k rows matter, because row i can't contain one of the k smallest pairs if i >= k (rows 0 to k - 1 each have a smaller-or-equal first pair).
Pop the smallest pair (i, j) and add it to the result. Push the next pair in its row, (i, j + 1).
Repeat until k pairs are taken or the heap is empty.

PriorityQueue is a min-heap (see heaps and PriorityQueue), so the smallest entry comes out first. A record is a small immutable class with generated accessors. The Entry record holds the sum and the two indexes, and the comparator orders entries by sum, then by i, then by j.

Java

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.PriorityQueue;

public class Main {
    record Entry(long sum, int i, int j) {}

    static List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        PriorityQueue<Entry> heap = new PriorityQueue<>(
            Comparator.comparingLong(Entry::sum).thenComparingInt(Entry::i).thenComparingInt(Entry::j));
        if (nums2.length > 0) {
            for (int i = 0; i < Math.min(k, nums1.length); i++) {
                heap.add(new Entry((long) nums1[i] + nums2[0], i, 0));
            }
        }
        List<int[]> result = new ArrayList<>();
        while (result.size() < k && !heap.isEmpty()) { // an empty heap means fewer than k pairs exist
            Entry e = heap.poll();
            result.add(new int[]{nums1[e.i()], nums2[e.j()]});
            if (e.j() + 1 < nums2.length) {
                heap.add(new Entry((long) nums1[e.i()] + nums2[e.j() + 1], e.i(), e.j() + 1));
            }
        }
        return result;
    }

    static String format(List<int[]> pairs) {
        StringBuilder sb = new StringBuilder("[");
        for (int i = 0; i < pairs.size(); i++) {
            if (i > 0) sb.append(", ");
            sb.append("(").append(pairs.get(i)[0]).append(", ").append(pairs.get(i)[1]).append(")");
        }
        return sb.append("]").toString();
    }

    public static void main(String[] args) {
        System.out.println(format(kSmallestPairs(new int[]{1, 7, 11}, new int[]{2, 4, 6}, 3)));
        System.out.println(format(kSmallestPairs(new int[]{1, 1, 2}, new int[]{1, 2, 3}, 2)));
        System.out.println(format(kSmallestPairs(new int[]{1, 2}, new int[]{3}, 3)));
        System.out.println(format(kSmallestPairs(new int[]{1, 2}, new int[]{}, 3)));
    }
}

Output

[(1, 2), (1, 4), (1, 6)]
[(1, 1), (1, 1)]
[(1, 3), (2, 3)]
[]

Time: O(k log k): the heap never holds more than k entries, and there are at most 2k pushes and k pops. Space: O(k).

Edge cases

k larger than m · n: every pair is returned; the heap runs empty.
An empty array: no pairs (the last example in the code).
Overflow: u + v can overflow int for values near the limits, and Java wraps silently. The code above casts to long before adding.