面试经典 150 题 2 —（堆）— 373. 查找和最小的 K 对数字

373. 查找和最小的 K 对数字

class Solution {
public:
    // priority_queue<参数一, 参数二, 参数三>：参数一是堆中节点的数据类型，参数二是承载底层数据结构堆（heap）的容器，参数三是对第一个参数的比较类
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        // 用 lambda 比较元素
        auto cmp = [&nums1, &nums2](const pair<int, int>& a, const pair<int, int>& b){
            return nums1[a.first] + nums2[a.second] > nums1[b.first] + nums2[b.second];
        };
        
        // 优先队列，大顶堆; decltype，在C++中，作为操作符，用于查询表达式的数据类型
        priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(cmp)> pri_queue(cmp);
        
        int nums1Size = nums1.size(), nums2Size = nums2.size();
        vector<vector<int>> answer;
        
        // 优先队列中存储的是下标对, 先将第一个数组的前k个元素的下标对放入优先队列， 第二个数组的下标为0，每次取出元素(x, y)时，只需要将第二个数组的下标加1即可，避免了重复加入元素
        for(int i = 0; i < min(k, nums1Size); i++){
            // pri_queue.emplace(i, 0);  // emplace()原地构造元素插入队列，并排序
            pri_queue.push(make_pair(i, 0)); 
        }
        
        // 取出前k个元素
        while(k-- > 0 && !pri_queue.empty()){
            // 取出堆顶元素
            auto [x, y] = pri_queue.top();
            pri_queue.pop();
            
            // 将取出的元素放入答案中
            answer.emplace_back(vector<int>{nums1[x], nums2[y]});
            
            // 如果第二个数组还有元素，将第一个数组的下标不变，第二个数组的下标加1，放入优先队列
            if(y + 1 < nums2Size){
                // pri_queue.emplace(x, y + 1);
                pri_queue.push(make_pair(x, y + 1));
            }
        }
        return answer;
    }
};