Find K Pairs with Smallest Sums

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.

Examples

Example 1

1
2
3


Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]]
Explanation: The first 3 pairs are returned from the sequence: [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

Example 2

1
2
3


Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [[1,1],[1,1]]
Explanation: The first 2 pairs are returned from the sequence: [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

Example 3

1
2
3


Input: nums1 = [1,2], nums2 = [3], k = 3
Output: [[1,3],[2,3]]
Explanation: All possible pairs are returned from the sequence: [1,3],[2,3]

Solution

Method 1 – Min Heap with BFS Expansion

We can use Kth Smallest Element in a Sorted Matrix#Method 4 - Using Priority Queue and Starting with First Row 🏆

Intuition

We use a min heap to always select the next smallest sum pair. By expanding only the next possible pairs, we avoid generating all possible pairs, making the solution efficient.

Approach

Initialize a min heap with pairs (nums1[i], nums2[0]) for i in range up to min(k, len(nums1)).
Pop the smallest sum pair from the heap and add it to the answer.
For each popped pair (i, j), push (i, j+1) into the heap if j+1 < len(nums2).
Repeat until k pairs are found or the heap is empty.
Handle edge cases where either array is empty or k is zero.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20


class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<vector<int>> ans;
        if (nums1.empty() || nums2.empty() || k == 0) return ans;
        auto cmp = [&](const tuple<int,int,int>& a, const tuple<int,int,int>& b) {
            return get<0>(a) > get<0>(b);
        };
        priority_queue<tuple<int,int,int>, vector<tuple<int,int,int>>, decltype(cmp)> pq(cmp);
        for (int i = 0; i < min((int)nums1.size(), k); ++i)
            pq.emplace(nums1[i]+nums2[0], i, 0);
        while (!pq.empty() && ans.size() < k) {
            auto [sum, i, j] = pq.top(); pq.pop();
            ans.push_back({nums1[i], nums2[j]});
            if (j+1 < nums2.size())
                pq.emplace(nums1[i]+nums2[j+1], i, j+1);
        }
        return ans;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19


func kSmallestPairs(nums1, nums2 []int, k int) [][]int {
    ans := [][]int{}
    if len(nums1) == 0 || len(nums2) == 0 || k == 0 {
        return ans
    }
    h := &minHeap{}
    heap.Init(h)
    for i := 0; i < min(len(nums1), k); i++ {
        heap.Push(h, pair{nums1[i]+nums2[0], i, 0})
    }
    for h.Len() > 0 && len(ans) < k {
        p := heap.Pop(h).(pair)
        ans = append(ans, []int{nums1[p.i], nums2[p.j]})
        if p.j+1 < len(nums2) {
            heap.Push(h, pair{nums1[p.i]+nums2[p.j+1], p.i, p.j+1})
        }
    }
    return ans
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16


class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> ans = new ArrayList<>();
        if (nums1.length == 0 || nums2.length == 0 || k == 0) return ans;
        PriorityQueue<int[]> pq = new PriorityQueue<>(Comparator.comparingInt(a -> a[0]));
        for (int i = 0; i < Math.min(nums1.length, k); i++)
            pq.offer(new int[]{nums1[i]+nums2[0], i, 0});
        while (!pq.isEmpty() && ans.size() < k) {
            int[] p = pq.poll();
            ans.add(Arrays.asList(nums1[p[1]], nums2[p[2]]));
            if (p[2]+1 < nums2.length)
                pq.offer(new int[]{nums1[p[1]]+nums2[p[2]+1], p[1], p[2]+1});
        }
        return ans;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14


class Solution {
    fun kSmallestPairs(nums1: List<Int>, nums2: List<Int>, k: Int): List<List<Int>> {
        val ans = mutableListOf<List<Int>>()
        if (nums1.isEmpty() || nums2.isEmpty() || k == 0) return ans
        val pq = PriorityQueue(compareBy<Pair<Int, Int>> { nums1[it.first] + nums2[it.second] })
        for (i in 0 until minOf(nums1.size, k)) pq.add(Pair(i, 0))
        while (pq.isNotEmpty() && ans.size < k) {
            val (i, j) = pq.poll()
            ans.add(listOf(nums1[i], nums2[j]))
            if (j+1 < nums2.size) pq.add(Pair(i, j+1))
        }
        return ans
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


def k_smallest_pairs(nums1: list[int], nums2: list[int], k: int) -> list[list[int]]:
    ans = []
    if not nums1 or not nums2 or k == 0:
        return ans
    heap = []
    for i in range(min(len(nums1), k)):
        heapq.heappush(heap, (nums1[i]+nums2[0], i, 0))
    while heap and len(ans) < k:
        _, i, j = heapq.heappop(heap)
        ans.append([nums1[i], nums2[j]])
        if j+1 < len(nums2):
            heapq.heappush(heap, (nums1[i]+nums2[j+1], i, j+1))
    return ans

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18


impl Solution {
    pub fn k_smallest_pairs(nums1: Vec<i32>, nums2: Vec<i32>, k: i32) -> Vec<Vec<i32>> {
        let mut ans = vec![];
        if nums1.is_empty() || nums2.is_empty() || k == 0 { return ans; }
        let mut heap = BinaryHeap::new();
        for i in 0..nums1.len().min(k as usize) {
            heap.push(Reverse((nums1[i]+nums2[0], i, 0)));
        }
        while let Some(Reverse((_, i, j))) = heap.pop() {
            ans.push(vec![nums1[i], nums2[j]]);
            if ans.len() == k as usize { break; }
            if j+1 < nums2.len() {
                heap.push(Reverse((nums1[i]+nums2[j+1], i, j+1)));
            }
        }
        ans
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21


class Solution {
    kSmallestPairs(nums1: number[], nums2: number[], k: number): number[][] {
        const ans: number[][] = [];
        if (!nums1.length || !nums2.length || k === 0) return ans;
        const heap: [number, number, number][] = [];
        for (let i = 0; i < Math.min(nums1.length, k); i++) {
            heap.push([nums1[i]+nums2[0], i, 0]);
        }
        heap.sort((a, b) => a[0] - b[0]);
        while (heap.length && ans.length < k) {
            const [_, i, j] = heap.shift()!;
            ans.push([nums1[i], nums2[j]]);
            if (j+1 < nums2.length) {
                let pair: [number, number, number] = [nums1[i]+nums2[j+1], i, j+1];
                heap.push(pair);
                heap.sort((a, b) => a[0] - b[0]);
            }
        }
        return ans;
    }
}

Complexity

⏰ Time complexity: O(k log k), each heap operation is log k and we do up to k operations.
🧺 Space complexity: O(k), storing up to k pairs in the heap and answer.

Problem#

Examples#

Example 1#

Example 2#

Example 3#

Solution#

Method 1 – Min Heap with BFS Expansion#

Intuition#

Approach#

Code#

Complexity#