Find K Pairs with Smallest Sums

Sanjeev SharmaSanjeev Sharma
2 min read

Advertisement

Problem

Given two sorted arrays, find k pairs (one from each) with the smallest sums.

Key insight: K-way merge. Initialize heap with pairs (nums1[i], nums2[0]) for i in [0, min(k, len1)). When popping (i, j), push (i, j+1).

Solutions

// C++
vector<vector<int>> kSmallestPairs(vector<int>& n1, vector<int>& n2, int k) {
    using T = tuple<int,int,int>; // sum, i, j
    priority_queue<T, vector<T>, greater<T>> minH;
    for (int i = 0; i < min(k, (int)n1.size()); i++)
        minH.push({n1[i]+n2[0], i, 0});
    vector<vector<int>> res;
    while (!minH.empty() && (int)res.size() < k) {
        auto [s, i, j] = minH.top(); minH.pop();
        res.push_back({n1[i], n2[j]});
        if (j + 1 < (int)n2.size()) minH.push({n1[i]+n2[j+1], i, j+1});
    }
    return res;
}
// Java
public List<List<Integer>> kSmallestPairs(int[] n1, int[] n2, int k) {
    PriorityQueue<int[]> minH = new PriorityQueue<>(Comparator.comparingInt(a -> n1[a[0]]+n2[a[1]]));
    for (int i = 0; i < Math.min(k, n1.length); i++) minH.offer(new int[]{i, 0});
    List<List<Integer>> res = new ArrayList<>();
    while (!minH.isEmpty() && res.size() < k) {
        int[] p = minH.poll();
        res.add(Arrays.asList(n1[p[0]], n2[p[1]]));
        if (p[1]+1 < n2.length) minH.offer(new int[]{p[0], p[1]+1});
    }
    return res;
}
// JavaScript
function kSmallestPairs(nums1, nums2, k) {
    const heap = [];
    for (let i = 0; i < Math.min(k, nums1.length); i++)
        heap.push([nums1[i] + nums2[0], i, 0]);
    heap.sort((a, b) => a[0] - b[0]);
    const result = [];
    while (heap.length && result.length < k) {
        const [, i, j] = heap.shift();
        result.push([nums1[i], nums2[j]]);
        if (j + 1 < nums2.length) {
            const entry = [nums1[i] + nums2[j+1], i, j+1];
            let pos = heap.findIndex(x => x[0] >= entry[0]);
            if (pos === -1) heap.push(entry);
            else heap.splice(pos, 0, entry);
        }
    }
    return result;
}
# Python
import heapq

def kSmallestPairs(nums1, nums2, k):
    if not nums1 or not nums2:
        return []
    heap = [(nums1[i] + nums2[0], i, 0) for i in range(min(k, len(nums1)))]
    heapq.heapify(heap)
    result = []
    while heap and len(result) < k:
        s, 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)
  • Space: O(k)

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro