Leetcode 373: Find K Pairs with Smallest Sums (MinHeap题)
- Find K Pairs with Smallest Sums
Medium
Companies
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 (u1, v1), (u2, v2), …, (uk, vk) with the smallest sums.
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]
Example 3:
Input: nums1 = [1,2], nums2 = [3], k = 3
Output: [[1,3],[2,3]]
Explanation: All possible pairs are returned from the sequence: [1,3],[2,3]
Constraints:
1 <= nums1.length, nums2.length <= 105
-109 <= nums1[i], nums2[i] <= 109
nums1 and nums2 both are sorted in non-decreasing order.
1 <= k <= 104
解法1：
用MinHeap. Create 自己的Node。注意一开始压栈的时候不要只压一个，而是对nums1所有的i, 压Node(nums1[i], nums2[0], 0)。这样以后的nums1就不用管了，只用管nums2的index往后移动即可。
struct Node {
int sum1;
int sum2;
int index2;
Node(int s1, int s2, int id2) : sum1(s1), sum2(s2), index2(id2) {}
bool operator < (const Node &n) const {
return sum1 + sum2 >= n.sum1 + n.sum2;
}
};
class Solution {
public:
vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
vector<vector<int>> res;
int n1 = nums1.size(), n2 = nums2.size();
priority_queue<Node> minHeap;
for (int i = 0; i < min(n1, k); i++) {
minHeap.push(Node(nums1[i], nums2[0], 0));
}
int count = 0;
while(!minHeap.empty() && count < k) {
Node topNode = minHeap.top();
minHeap.pop();
res.push_back({topNode.sum1, topNode.sum2});
if (topNode.index2 < n2 - 1) {
minHeap.push(Node(topNode.sum1, nums2[topNode.index2 + 1], topNode.index2 + 1));
}
count++;
}
return res;
}
};