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

题目链接

Leetcode.1439 有序矩阵中的第 k 个最小数组和 Rating : 2134

题目描述

给你一个 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

提示：

• $m==mat.length$
• $n==mat.length[i]$
• $1<=m,n<=40$
• $1<=k<=min(200,n^m)$
• $1<=mat[i][j]<=5000$
• $mat[i]$ 是一个非递减数组

解法一：暴力

假设我们已经把 前 $i-1$ 行的 $k$ 最小数组和记录在了一个数组 $a$ 中，那么我们只需要把这个 $a$ 中元素 和 第$i$行的元素 两两相加，再取前 $k$ 个最小的元素。一种重复这个过程，那么最后$a[k-1]$ 就是我们要求的值。

时间复杂度： $O(m\times n\times k\times log(nk))$

C++代码：

class Solution {
public:
    int kthSmallest(vector<vector<int>>& mat, int k) {
        vector<int> a{0};

        for(auto &e:mat){
            vector<int> b;
            for(auto &x:a){
                for(auto &y:e) b.push_back(x + y);
            }
            sort(b.begin(),b.end());
            if(b.size() > k) b.resize(k);
            a = move(b);
        }
        return a.back();
    }
};

解法二：堆

参考这篇题解

时间复杂度： $O(m\times k\times logmin(n,k))$

C++代码：

class Solution {
public:
    int kthSmallest(vector<vector<int>>& mat, int k) {
        vector<int> ans{0};

        auto fun = [&](vector<int>& a,vector<int>& b){
            int m = a.size() , n = b.size();
            priority_queue<tuple<int,int,int>> q;
            q.emplace(-a[0]-b[0],0,0);

            vector<int> ans;
            while(!q.empty() && ans.size() < k){
                auto [_,i,j] = q.top();
                q.pop();

                ans.push_back(a[i] + b[j]);
                if(j == 0 && i + 1 < m) q.emplace(-a[i + 1] - b[j],i + 1,j);
                if(j + 1 < n) q.emplace(-a[i] - b[j + 1] , i , j + 1);
            }
            return ans;
        };

        for(auto &row:mat){
            ans = fun(ans,row);
        }
        return ans.back();
    }
};