有序矩阵中的第 k 个最小数组和-小顶堆法

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

题目描述

方法一

从上到下遍历矩阵的所有行，假设计算出了前 $i-1$ 行形成的前 $k$ 个最小数组和（记作 $sum$ ），遍历到第 $i$ 行时，把 $sum$ 与第 $i$ 行的数两两相加，然后只保留其中最小的 $k$ 个数，作为新的 $sum$ ，然后继续遍历矩阵的下一行，重复该过程直至结束。
代码如下:

class Solution {
public:
    int kthSmallest(vector<vector<int>> &mat, int k) {
        vector<int> a = {0};
        for (auto &row: mat) {
            vector<int> b;
            for (int x: a)
                for (int y: row)
                    b.push_back(x + y);
            sort(b.begin(), b.end());
            // 保留最小的 k 个
	    if (b.size() > k) 
            	b.resize(k);
            a = move(b);
        }
        return a.back();
    }
};

调用std::move函数，该函数是将一个左值转换为右值，之后便是利用移动构造函数进行拷贝，能够提高效率。
时间复杂度为 O(mnklogk), m是矩阵行数，n是列数。

方法二

其中，$l_g$ 个序列的首项为: $$f[0]+g[0],f[0]+g[1],...,f[0]+g[l_g-1]$$
将首项存入堆中之后，可以访问到它的首个元素，也就是最小的元素。此后，将该元素的下一项存入堆中。如此反复k次，即可求得答案。通过阅读官方题解我们也可以发现，方法一的暴力法效率不高的原因在于遍历了矩阵的每一行的每一个元素，而事实上融合优先队列，只需要k次处理即可。
代码如下：

class Solution {
public:
    int kthSmallest(vector<vector<int>>& mat, int k) {
        int m = mat.size();
        vector<int> prev = mat[0];
        for (int i = 1; i < m; ++i) 
        {
            //移动构造函数，提高性能
            prev = move(merge(prev, mat[i], k));
        }
        return prev[k - 1];
    }

    vector<int> merge(const vector<int>& f, const vector<int>& g, int k) {
        if (g.size() > f.size()) {
            return merge(g, f, k);
        }

        priority_queue<Entry> q;
        for (int i = 0; i < g.size(); ++i) {
            q.emplace(0, i, f[0] + g[i]);
        }

        vector<int> ans;
        while (k && !q.empty()) {
            Entry entry = q.top();
            q.pop();
            ans.push_back(entry.sum);
            if (entry.i + 1 < f.size()) 
            {
                q.emplace(entry.i + 1, entry.j, f[entry.i + 1] + g[entry.j]);
            }
            --k;
        }

        return ans;
    }

private:
    struct Entry {
        int i, j, sum;
        Entry(int _i, int _j, int _sum): i(_i), j(_j), sum(_sum) {}
        //与sort的cmp不同，operator<为真，就把sum所在的对象排在后边
        bool operator< (const Entry& that) const {
            return sum > that.sum;
        }
    };
};