Medium
Find K Pairs with Smallest Sums — C++
Full explanation · Time O(k * log(min(n, m, k))) · Space O(min(n, m, k))
// Time: O(k * log(min(n, m, k))), where n is the size of num1, and m is the size of num2.
// Space: O(min(n, m, k))
class Solution {
public:
vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
vector<pair<int, int>> pairs;
if (nums1.size() > nums2.size()) {
vector<pair<int, int>> tmp = kSmallestPairs(nums2, nums1, k);
for (const auto& pair : tmp) {
pairs.emplace_back(pair.second, pair.first);
}
return pairs;
}
using P = pair<int, pair<int, int>>;
priority_queue<P, vector<P>, greater<P>> q;
auto push = [&nums1, &nums2, &q](int i, int j) {
if (i < nums1.size() && j < nums2.size()) {
q.emplace(nums1[i] + nums2[j], make_pair(i, j));
}
};
push(0, 0);
while (!q.empty() && pairs.size() < k) {
auto tmp = q.top(); q.pop();
int i, j;
tie(i, j) = tmp.second;
pairs.emplace_back(nums1[i], nums2[j]);
push(i, j + 1);
if (j == 0) {
push(i + 1, 0); // at most queue min(m, n) space.
}
}
return pairs;
}
};