Find K Pairs with Smallest Sums

medium heap

Problem

You are given two integer arrays nums1 and nums2 sorted in ascending order and an integer k. Define a pair (u, v) which consists of one element from the first array and one element from the second array. Return the k pairs (u1, v1), (u2, v2), ..., (uk, vk) with the smallest sums.

Inputnums1 = [1,7,11], nums2 = [2,4,6], k = 3

Output[[1,2],[1,4],[1,6]]

Seed a min-heap with (nums1[0]+nums2[0], 0, 0). On each pop (i, j), push (i, j+1); also push (i+1, 0) the first time j is 0. The heap holds at most k frontier pairs.

nums1

nums2

import heapq

def k_smallest_pairs(nums1, nums2, k):
    if not nums1 or not nums2: return []
    heap = [(nums1[0] + nums2[0], 0, 0)]
    seen = {(0, 0)}
    out = []
    while heap and len(out) < k:
        s, i, j = heapq.heappop(heap)
        out.append([nums1[i], nums2[j]])
        for di, dj in ((0, 1), (1, 0)):
            ni, nj = i + di, j + dj
            if ni < len(nums1) and nj < len(nums2) and (ni, nj) not in seen:
                seen.add((ni, nj))
                heapq.heappush(heap, (nums1[ni] + nums2[nj], ni, nj))
    return out

// Uses a binary min-heap on [sum, i, j] (heap class omitted).
function kSmallestPairs(nums1, nums2, k) {
  if (!nums1.length || !nums2.length) return [];
  const heap = new MinHeap((a, b) => a[0] - b[0]);
  const seen = new Set(["0,0"]);
  heap.push([nums1[0] + nums2[0], 0, 0]);
  const out = [];
  while (heap.size() && out.length < k) {
    const [, i, j] = heap.pop();
    out.push([nums1[i], nums2[j]]);
    for (const [di, dj] of [[0, 1], [1, 0]]) {
      const ni = i + di, nj = j + dj, key = ni + "," + nj;
      if (ni < nums1.length && nj < nums2.length && !seen.has(key)) {
        seen.add(key);
        heap.push([nums1[ni] + nums2[nj], ni, nj]);
      }
    }
  }
  return out;
}

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] a, int[] b, int k) {
        List<List<Integer>> out = new ArrayList<>();
        if (a.length == 0 || b.length == 0) return out;
        PriorityQueue<int[]> pq = new PriorityQueue<>((x, y) -> x[0] - y[0]);
        Set<Long> seen = new HashSet<>();
        pq.add(new int[]{a[0] + b[0], 0, 0});
        seen.add(0L);
        while (!pq.isEmpty() && out.size() < k) {
            int[] t = pq.poll();
            int i = t[1], j = t[2];
            out.add(Arrays.asList(a[i], b[j]));
            int[][] nb = {{i, j + 1}, {i + 1, j}};
            for (int[] p : nb) {
                if (p[0] < a.length && p[1] < b.length) {
                    long key = (long) p[0] * 100000 + p[1];
                    if (seen.add(key)) pq.add(new int[]{a[p[0]] + b[p[1]], p[0], p[1]});
                }
            }
        }
        return out;
    }
}

vector<vector<int>> kSmallestPairs(vector<int>& a, vector<int>& b, int k) {
    vector<vector<int>> out;
    if (a.empty() || b.empty()) return out;
    priority_queue<tuple<int,int,int>, vector<tuple<int,int,int>>, greater<>> pq;
    set<pair<int,int>> seen;
    pq.emplace(a[0] + b[0], 0, 0);
    seen.insert({0, 0});
    while (!pq.empty() && (int) out.size() < k) {
        auto [s, i, j] = pq.top(); pq.pop();
        out.push_back({a[i], b[j]});
        for (auto [di, dj] : {make_pair(0,1), make_pair(1,0)}) {
            int ni = i + di, nj = j + dj;
            if (ni < (int) a.size() && nj < (int) b.size() && seen.insert({ni, nj}).second)
                pq.emplace(a[ni] + b[nj], ni, nj);
        }
    }
    return out;
}

Explanation

There can be up to len(nums1) * len(nums2) pairs, but we only need the k smallest sums. The key insight is that because both arrays are sorted, the smallest pair is always (nums1[0], nums2[0]), and the next smallest must be a neighbor of a pair we already took. A min-heap lets us always grab the current smallest cheaply.

We track pairs by their index positions (i, j). We seed the heap with the sum and indices of the corner pair (0, 0), and a seen set so we never add the same index pair twice.

Each round we pop the smallest-sum entry, record the actual pair [nums1[i], nums2[j]] into out, then push its two neighbors: (i, j+1) (next element of the second array) and (i+1, j) (next element of the first array). The heap stays small because we only ever expand the frontier of pairs we've reached.

We stop once we've collected k pairs (or the heap empties). Since the heap always hands back the smallest available sum, the pairs come out in non-decreasing sum order.

Example: nums1 = [1,7,11], nums2 = [2,4,6], k = 3. Pop (1,2) sum 3, push neighbors; pop (1,4) sum 5; pop (1,6) sum 7. Result: [[1,2],[1,4],[1,6]].

Time: O(k log k) Space: O(k)