力扣 1439. 有序矩阵中的第 k 个最小数组和 [堆]

1439. 有序矩阵中的第 k 个最小数组和 #

给你一个 m * n 的矩阵 mat，以及一个整数 k ，矩阵中的每一行都以非递减的顺序排列。

你可以从每一行中选出 1 个元素形成一个数组。返回所有可能数组中的第 k 个最小数组和。

示例 1：

输入：mat = [[1,3,11],[2,4,6]], k = 5
输出：7
解释：从每一行中选出一个元素，前 k 个和最小的数组分别是：
[1,2], [1,4], [3,2], [3,4], [1,6]。其中第 5 个的和是 7 。

示例 2：

输入：mat = [[1,3,11],[2,4,6]], k = 9
输出：17

示例 3：

输入：mat = [[1,10,10],[1,4,5],[2,3,6]], k = 7
输出：9
解释：从每一行中选出一个元素，前 k 个和最小的数组分别是：
[1,1,2], [1,1,3], [1,4,2], [1,4,3], [1,1,6], [1,5,2], [1,5,3]。其中第 7 个的和是 9 。

示例 4：

输入：mat = [[1,1,10],[2,2,9]], k = 7
输出：12

题解#

此题和很力扣 373. 查找和最小的 K 对数字 [堆]相似，建议先做373，然后可以直接把kSmallestPairs函数复制过来，修改加以利用。

此题和373的区别就是 373是只有两个数组nums1和nums2（长度可能不等），每个数组选出一个，组成值对u,v，比较和，而且返回的也是值对。而此题是有m个数组，每个长n，每个数组都要选出一个元素，所以我们通过373中求两个数组nums1和nums2的函数，修改返回值为其和，而不是值对。

与373的逻辑一样，利用非递减的性质，mat[0]即为nums1，那么mat[1]为num2，传入kSmallestPairs，得到的结果ans为第一行与第二行中和为前k小的结果，然后将ans当作nums1,mat[2]作为nums2，再次传入，直到mat[m-1]。

通过每次处理两个数组，完成题目要求的处理m个数组的要求。

查看代码

 class Solution {
public:
    vector<int> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        //排序方法：若前一个pair的两个下标得到的元素之和>第二个pair得到的元素之和，顺序不变，否则交换,即按和的从大到小排列，得到小顶堆，top为最小值
        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];
        };

        int n1 = nums1.size();
        int n2 = nums2.size();
        vector<int> ans;   
        priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(cmp)> pq(cmp);
        for (int i = 0; i < min(k,n1); i++) {//num1中每个元素分别与num2中第一个元素组成下标对
            pq.emplace(i, 0);
        }
        while (k-- > 0 && !pq.empty()) {
            auto [x, y] = pq.top(); //按和从小到大顺序获取下标对，假如当前为第m小的
            pq.pop();
            //这里保存当前数组的和
            ans.emplace_back(nums1[x]+nums2[y]);
            //接下来num2的下标右移，入堆，堆会自动排序，这样下一轮循环的时候，top会得到m+1小的
            if (y + 1 < n2) {
                pq.emplace(x, y + 1);
            }
        }
        return ans;
    }
    
    int kthSmallest(vector<vector<int>>& mat, int k) {
        vector<int> ans=mat[0];//第一行
        int m=mat.size();
        for(int i=1;i<m;i++){
            ans=kSmallestPairs(ans,mat[i],k);
        }
        return ans.back();
    }
};