LC1439 有序矩阵中的第 k 个最小数组和

题目要求返回所有可能数组中的第 k 个 最小 数组和。

根据题目给出的矩阵每行均为单调递增的提示,可以得出下一个最小数组和是由之前得到过的数组转移过来的。

因为k的范围不大,所以可以使用优先队列+bfs思想模拟这个过程。

 

备忘这里的优先队列重载方法

class Solution {
public:
    int kthSmallest(vector<vector<int>>& mat, int k) {
        struct node {
            vector<int> pos;
            int tot;
            node(vector<int> tmp, int t) {
                pos = tmp, tot = t;
            }
        };
        auto my_cmp = [](const node &a, const node &b) {
            return a.tot > b.tot;
        };
        priority_queue<node, vector<node>, decltype(my_cmp)> pq(my_cmp);

        int M = mat.size(), N = mat[0].size();
        vector<int> tmp(M);
        set<vector<int> > s;
        int sum = 0;
        for(int i = 0; i < M; i++) {
            tmp[i] = 0, sum += mat[i][0];
        }
        pq.push({tmp, sum});
        while(!pq.empty()) {
            tmp = pq.top().pos, sum = pq.top().tot;
            pq.pop();
            k--;
            /*printf("%d\n", sum);
            for(int i = 0; i < M; i++)
                printf("%d ", tmp[i]);
            puts("");*/
            if(k == 0) break;
            for(int i = 0; i < M; i++) {
                if(tmp[i]+1 < N) {
                    tmp[i]++;
                    if(s.count(tmp) == 0) {
                        s.insert(tmp);
                        pq.push({tmp, sum+mat[i][tmp[i]]-mat[i][tmp[i]-1]});
                    }
                    tmp[i]--;
                }
            }
        }
        return sum;
    }
};
View Code

 

posted @ 2020-08-03 16:36  参婵  阅读(133)  评论(0编辑  收藏  举报