Given two sorted integer arrays nums1
and nums2
, and an integer k
, your task is to find the k
pairs with the smallest sums. Each pair consists of one element from nums1
and one element from nums2
. The result should be a list of pairs [u, v]
such that u
is from nums1
and v
is from nums2
, and their sum is among the k
smallest possible sums.
Key constraints:
k
pairs, or fewer if there aren't enough possible pairs.
At first glance, you might consider generating all possible pairs between nums1
and nums2
, compute their sums, and then sort them to get the k
smallest. This brute-force approach is simple but inefficient, especially when the arrays are large, since the number of pairs can be huge (len(nums1) * len(nums2)
).
To optimize, consider that both arrays are sorted. This means the smallest sums will always involve the smallest elements from each array. We can use a min-heap (priority queue) to efficiently keep track of the next smallest possible pair sum, and only generate new pairs as needed, rather than all at once.
The challenge is to systematically explore pairs in order of increasing sum, without missing any, and without adding duplicates.
We'll use a min-heap to always extract the next smallest sum pair efficiently. Here's the step-by-step process:
nums1
with every element of nums2
up to k
(since the arrays are sorted, these are the smallest possible sums for each nums2
element).(sum, i, j)
into the heap, where i
is the index in nums1
, and j
is the index in nums2
.nums1
(i.e., (i+1, j)
) into the heap. This ensures we explore all combinations in sorted order, but never more than necessary.k
pairs or the heap is empty.(i+1, j)
for each pop, not all combinations, ensuring no duplicate pairs are generated.This approach leverages the sorted property of the arrays and the efficiency of the heap to avoid unnecessary work.
Let's work through an example:
Input: nums1 = [1, 7, 11]
, nums2 = [2, 4, 6]
, k = 3
(1+2, 0, 0) = (3, 0, 0)
, (1+4, 0, 1) = (5, 0, 1)
, (1+6, 0, 2) = (7, 0, 2)
Notice how we only ever push the next candidate for each nums1
index, ensuring no duplicates and minimal heap size.
Brute-force approach:
O(mn \log(mn))
to generate all pairs and sort, where m = len(nums1)
, n = len(nums2)
.O(mn)
to store all pairs.O(k \log k)
because we push/pop at most k
times from the heap and the heap size is at most k
.O(k)
for the heap and result list.
The optimized approach is much more efficient, especially for large arrays and small k
.
By leveraging the sorted nature of the input arrays and a min-heap, we efficiently find the k
smallest sum pairs without generating all possible pairs. The key insight is to always generate the next smallest candidate using the heap and only push necessary pairs, ensuring optimal time and space usage. This approach is elegant, scalable, and avoids brute-force inefficiency.
import heapq
class Solution:
def kSmallestPairs(self, nums1, nums2, k):
if not nums1 or not nums2 or k == 0:
return []
min_heap = []
res = []
# Only need first k elements from nums2 since arrays are sorted
for j in range(min(k, len(nums2))):
heapq.heappush(min_heap, (nums1[0] + nums2[j], 0, j))
while min_heap and len(res) < k:
curr_sum, i, j = heapq.heappop(min_heap)
res.append([nums1[i], nums2[j]])
if i + 1 < len(nums1):
heapq.heappush(min_heap, (nums1[i + 1] + nums2[j], i + 1, j))
return res
#include <vector>
#include <queue>
using namespace std;
class Solution {
public:
vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
vector<vector<int>> res;
if (nums1.empty() || nums2.empty() || k == 0) return res;
auto cmp = [](const vector<int>& a, const vector<int>& b) {
return a[0] + a[1] > b[0] + b[1];
};
priority_queue<vector<int>, vector<vector<int>>, decltype(cmp)> minHeap(cmp);
for (int j = 0; j < min(k, (int)nums2.size()); ++j) {
minHeap.push({nums1[0], nums2[j], 0});
}
while (!minHeap.empty() && res.size() < k) {
auto curr = minHeap.top(); minHeap.pop();
int n1 = curr[0], n2 = curr[1], i = curr[2];
res.push_back({n1, n2});
if (i + 1 < nums1.size()) {
minHeap.push({nums1[i + 1], n2, i + 1});
}
}
return res;
}
};
import java.util.*;
class Solution {
public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
List<List<Integer>> res = new ArrayList<>();
if (nums1.length == 0 || nums2.length == 0 || k == 0) return res;
PriorityQueue<int[]> minHeap = new PriorityQueue<>((a, b) -> (a[0] + a[1]) - (b[0] + b[1]));
for (int j = 0; j < Math.min(k, nums2.length); j++) {
minHeap.offer(new int[]{nums1[0], nums2[j], 0});
}
while (!minHeap.isEmpty() && res.size() < k) {
int[] curr = minHeap.poll();
res.add(Arrays.asList(curr[0], curr[1]));
int i = curr[2];
if (i + 1 < nums1.length) {
minHeap.offer(new int[]{nums1[i + 1], curr[1], i + 1});
}
}
return res;
}
}
// MinHeap implementation for JS
class MinHeap {
constructor() { this.heap = []; }
push(val) {
this.heap.push(val);
let i = this.heap.length - 1;
while (i > 0) {
let p = Math.floor((i - 1) / 2);
if (this.heap[p][0] <= this.heap[i][0]) break;
[this.heap[p], this.heap[i]] = [this.heap[i], this.heap[p]];
i = p;
}
}
pop() {
if (this.heap.length === 1) return this.heap.pop();
const ret = this.heap[0];
this.heap[0] = this.heap.pop();
let i = 0, n = this.heap.length;
while (true) {
let l = 2 * i + 1, r = 2 * i + 2, min = i;
if (l < n && this.heap[l][0] < this.heap[min][0]) min = l;
if (r < n && this.heap[r][0] < this.heap[min][0]) min = r;
if (min === i) break;
[this.heap[i], this.heap[min]] = [this.heap[min], this.heap[i]];
i = min;
}
return ret;
}
size() { return this.heap.length; }
}
var kSmallestPairs = function(nums1, nums2, k) {
if (!nums1.length || !nums2.length || k === 0) return [];
let minHeap = new MinHeap();
let res = [];
for (let j = 0; j < Math.min(k, nums2.length); j++) {
minHeap.push([nums1[0] + nums2[j], 0, j]);
}
while (minHeap.size() && res.length < k) {
let [sum, i, j] = minHeap.pop();
res.push([nums1[i], nums2[j]]);
if (i + 1 < nums1.length) {
minHeap.push([nums1[i + 1] + nums2[j], i + 1, j]);
}
}
return res;
};