堆的运用——有序元素的多路归并topk问题

Posted 2022-02-06 C+++++++++++++++++++

文章目录

• • 题目一：有序矩阵第k小的元素(提炼出做题方法)
  • 解题技法
  • 解题代码
  • (进阶运用)题目二：查找和最小的K对数字
  • 题目解析
  • 解题代码

题目一：有序矩阵第k小的元素(提炼出做题方法)

题目链接

解题技法

感觉这张图基本就清楚了这题目如何解。

• 具体详解过程请看lc大神：题目详解

解题代码

class Solution 
public:
    //TODO 多路归并
    int kthSmallest(vector<vector<int>>& matrix, int k) 
        auto cmp = [&](pair<int,int>&a,pair<int,int>&b)
            return matrix[a.first][a.second]>matrix[b.first][b.second];
        ;
        priority_queue<pair<int,int>,vector<pair<int,int>>,decltype(cmp)>pq(cmp);
        int n = matrix.size();
        for(int i=0;i<min(k,n);i++)
            pq.push(i,0);//TODO 得到第一次的行首元素
        
        int ret = INT_MAX;
        while(k--&&!pq.empty())
            auto [x,y] = pq.top();pq.pop();
            if(y+1<n)//TODO 更新这一行的下一个元素到堆中
                pq.push(x,y+1);
            ret = matrix[x][y];
        
        return ret;
    
;

(进阶运用)题目二：查找和最小的K对数字

题目链接

题目解析

• 和前面那道题的做法一样，这道题是由于者均有序，所以如果是直接进行两层循环的枚举的话，得到的数字可以看作是一个和上题一模一样的矩阵，也就是把 nums1[0...]+nums2[0] 看作是一行的首元素即可，然后处理过程就和前面的处理过程是完全一致。和前面一题的区别仅仅在于未有确定矩阵的内容而已，而我们需要做的就是确定这个矩阵的内容！

细节优化：由于矩阵的内容由我们来确定，为了防止初始化矩阵首行元素过多，我们可以采取把长度小的 nums 作为行的标准，那么为了让每次的答案顺序不变，所以需要一个标记。

解题代码

class Solution 
public:
    bool flag = true;
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) 
        vector<vector<int>> ans;
        int n = nums1.size(), m = nums2.size();
        if(n > m)  //始终确保nums1为两数组中长度较少的那个(这样做可以适当的减少堆的初始大小)，这个不处理也可以，只是简单的优化
            swap(nums1, nums2);
            swap(m,n);
            flag = false;//确保原本的第一个取数的数字时nums1原本的数字
        
        //定义比较规则
        auto cmp = [&](const auto& a, const auto& b)
            return nums1[a.first] + nums2[a.second] > nums1[b.first] + nums2[b.second];
        ;
        priority_queue< pair<int,int>, vector<pair<int,int>>, decltype(cmp) > q(cmp);
        for(int i = 0; i < min(n,k); i++)
            q.push( i, 0 );
        
        while(k-- and !q.empty())
            auto [a,b] = q.top();
            q.pop();
            flag ? ans.push_back( nums1[a], nums2[b]) : ans.push_back( nums2[b], nums1[a]);            
            //TODO 得到这一行的下一个元素，如果超过则不入
            if(b + 1 < m) q.push( a, b + 1 );
        
        return ans;
    
;