Find K Pairs with Smallest Sums
Medium · Heap / Priority Queue · Explanation and code example
Heap Merge of Pair-Sum Rows
How it works
For a fixed first-array index, pairing it with successive second-array values creates a sorted row of sums. Seed a min-heap with the first pair from each relevant row. After removing a pair, advance only that row to merge the globally smallest pairs.
Steps
- Seed up to min(k, len(nums1)) rows at second-array index zero.
- Pop the minimum sum and append its values.
- Push the next pair in the same row until k pairs or heap exhaustion.
Complexity
- Time: O((k + r) log(r + 1)), with r = min(k, len(nums1)).
- Space: O(r) heap work plus O(k) returned pairs.
Things to watch for
- Inputs must be sorted.
- Return pairs of values, rather than their indices.
- Duplicate values can form distinct valid index pairs.
Python code example
import heapq
class Solution:
def kSmallestPairs(self, nums1: list[int],
nums2: list[int],
k: int) -> list[list[int]]:
minHeap = []
for i in range(min(k, len(nums1))):
heapq.heappush(minHeap, (nums1[i] + nums2[0], i, 0))
ans = []
while minHeap and len(ans) < k:
_, i, j = heapq.heappop(minHeap)
ans.append([nums1[i], nums2[j]])
if j + 1 < len(nums2):
heapq.heappush(minHeap, (nums1[i] + nums2[j + 1], i, j + 1))
return ans
C++ code example
struct T {
int i;
int j;
int sum; // nums1[i] + nums2[j]
};
class Solution {
public:
vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2,
int k) {
vector<vector<int>> ans;
auto compare = [&](const T& a, const T& b) { return a.sum > b.sum; };
priority_queue<T, vector<T>, decltype(compare)> minHeap(compare);
for (int i = 0; i < k && i < nums1.size(); ++i)
minHeap.emplace(i, 0, nums1[i] + nums2[0]);
while (!minHeap.empty() && ans.size() < k) {
const auto [i, j, _] = minHeap.top();
minHeap.pop();
ans.push_back({nums1[i], nums2[j]});
if (j + 1 < nums2.size())
minHeap.emplace(i, j + 1, nums1[i] + nums2[j + 1]);
}
return ans;
}
};
Source and license
Explanation by LC Tracker. Code adapted from walkccc / Peng-Yu Chen · MIT license. Node types are supplied by LeetCode where needed.