k Smallest Sum Pairs From Two Arrays

medium heap

Problem

Given two ascending arrays nums1 and nums2 and an integer k, return the k pairs (u, v) with the smallest values of u + v, where u comes from nums1 and v from nums2.

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;
}

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