373. Find K Pairs with Smallest Sums
1. Description
You are given two integer arrays nums1 and nums2 sorted in non-decreasing 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 ($u_1$, $v_1$), ($u_2$, $v_2$), …, ($u_k$, $v_k$) with the smallest sums.
2. Example
Example 1
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
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]
3. Constraints
- 1 <= nums1.length, nums2.length <= 10$^5$
- -10$^9$ <= nums1[i], nums2[i] <= 10$^9$
- nums1 and nums2 both are sorted in non-decreasing order.
- 1 <= k <= 10$^4$
- k <= nums1.length * nums2.length
4. Solutions
Greedy && Heap
m = nums1.size(), n = nums2.size()
Time complexity: O(klogmin(m, n, k))
Space complexity: O(min(m, n, k))
class Solution {
public:
vector<vector<int>> kSmallestPairs(const vector<int> &nums1, const vector<int> &nums2, int k) {
bool nums1_smaller = nums1.size() < nums2.size();
const vector<int> &array1 = nums1_smaller ? nums2 : nums1;
const vector<int> &array2 = nums1_smaller ? nums1 : nums2;
const int m = array1.size(), n = array2.size();
auto compare = [&array1, &array2](const pair<int, int> &a, const pair<int, int> &b) {
return array1[a.first] + array2[a.second] > array1[b.first] + array2[b.second];
};
priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(compare)> indexes(compare);
for (int i = 0; i < min(n, k); ++i) {
indexes.push({0, i});
}
vector<vector<int>> results;
results.reserve(k);
for (int i = 0; i < k && !indexes.empty(); ++i) {
auto [row, column] = indexes.top();
indexes.pop();
results.push_back(
nums1_smaller ? vector<int>{array2[column], array1[row]}
: vector<int>{array1[row], array2[column]});
if (row + 1 < m) {
indexes.push({row + 1, column});
}
}
return results;
}
};