有序矩阵中的第 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

说明：

m ，n为矩阵 mat的行列数

1 <= m, n <= 40

1 <= k <= min(200, n ^ m)

1 <= mat[i][j] <= 5000

mat[i] 是一个非递减数组

可参考以下main函数：

int main()

{

    int m, n,data,k;

    vector<vector<int> > mat;

    cin>>m>>n;

    for(int i=0; i<m; i++)

    {

        vector<int> row;

        for(int j=0; j<n; j++)

        {

            cin>>data;

            row.push_back(data);

        }

        mat.push_back(row);

    }

    cin>>k;

    int res=Solution().kthSmallest(mat,k);

    cout<<res<<endl;

    return 0;

}

输入说明

首先输入矩阵mat的行列数m和n

然后输入m行，每行n个非负整数，表示mat的元素值，以空格分隔。

最后输入k。

输出说明

输出一个整数，表示结果。

输入范例

2 3
1 3 11
2 4 6
5

输出范例

7
 

实现代码

#include<iostream>

#include<queue>

#include<vector>

#include<algorithm>

#include<set>

using namespace std;



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

        int row = mat.size();//行
        int col = mat[0].size();//列

        vector<int> index(row);//创建一个大小为row的vector，每个默认值为0

        int firstnum = 0;
        for(int i = 0;i<row;i++){
            firstnum += mat[i][0];//计算第一列所有元素的和
        }

        set<vector<int>> visited;
        visited.insert(index);//记录访问过的元素下标

        priority_queue <pair<int,vector<int>>,vector<pair<int,vector<int>>>,greater<>> pq;
        pair<int,vector<int>> p{firstnum,index};
        pq.push(p);

        while(--k){
            pair<int,vector<int>> cur = pq.top();
            pq.pop();
            for(int i = 0;i<row;i++){
                if(cur.second[i] < col-1){
                    vector<int> v(cur.second);
                    v[i]++;
                    if(!visited.count(v)){
                        visited.insert(v);
                        int nextnum = cur.first - mat[i][cur.second[i]] + mat[i][v[i]];
                        pq.push(make_pair(nextnum,v));
                    }
                }
            }
        }

        return pq.top().first;


    }
};

int main()

{

    int m, n,data,k;

    vector<vector<int> > mat;

    cin>>m>>n;

    for(int i=0; i<m; i++)

    {

        vector<int> row;

        for(int j=0; j<n; j++)

        {

            cin>>data;

            row.push_back(data);

        }

        mat.push_back(row);

    }

    cin>>k;

    int res=Solution().kthSmallest(mat,k);

    cout<<res<<endl;

    return 0;

}