Computer Science/C++
[C++] priority queue (pair<int, pair<int, int>>), minheap
focalpoint
2023. 9. 2. 03:50
LeetCode: 373. Find K Pairs with Smallest Sums
Find K Pairs with Smallest Sums - LeetCode
Can you solve this real interview question? Find K Pairs with Smallest Sums - 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
leetcode.com
class Solution {
public:
vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
int m = nums1.size();
int n = nums2.size();
vector<vector<int>> ans;
set<pair<int, int>> visited;
priority_queue<pair<int, pair<int, int>>, vector<pair<int, pair<int, int>>>,
greater<pair<int, pair<int, int>>>> minHeap;
minHeap.push({nums1[0] + nums2[0], {0, 0}});
visited.insert({0, 0});
while (k-- && !minHeap.empty()) {
auto top = minHeap.top();
minHeap.pop();
int i = top.second.first;
int j = top.second.second;
ans.push_back({nums1[i], nums2[j]});
if (i + 1 < m && visited.find({i + 1, j}) == visited.end()) {
minHeap.push({nums1[i + 1] + nums2[j], {i + 1, j}});
visited.insert({i + 1, j});
}
if (j + 1 < n && visited.find({i, j + 1}) == visited.end()) {
minHeap.push({nums1[i] + nums2[j + 1], {i, j + 1}});
visited.insert({i, j + 1});
}
}
return ans;
}
};