Find K Pairs with Smallest Sums
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
← Previous
Ugly Number II