【LeetCode-378】有序矩阵中第 K 小的元素

问题

给你一个 n x n 矩阵 matrix ，其中每行和每列元素均按升序排序，找到矩阵中第 k 小的元素。
请注意，它是 排序后 的第 k 小元素，而不是第 k 个 不同 的元素。

示例

输入： matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
输出： 13
解释： 矩阵中的元素为 [1,5,9,10,11,12,13,13,15]，第 8 小元素是 13

解答1：优先队列

class Solution {
public:
    typedef tuple<int, int, int> TI;
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        priority_queue<TI, vector<TI>, greater<TI>> pq;
        auto push = [&](int i, int j) {
            if (i < matrix.size() && j < matrix.size())
                pq.emplace(matrix[i][j], i, j);
        };
        push(0, 0);
        int res = 0;
        for (; k > 0; k--) {
            auto [s, i, j] = pq.top(); pq.pop();
            res = s;
            push(i, j + 1);
            if (j == 0) push(i + 1, j);
        }
        return res;
    }
};

重点思路

利用了每行都是有序数列，并且第一列是有序数列这两个条件。但是其他列也是有序数列这个条件没用上。这个方法与【LeetCode-373】找出第 k 小的距离对相似，时间复杂度为\(O(klog(n))\)，最坏为\(O(n^2log(n))\)。

解答2：二分答案

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int n = matrix.size(), left = matrix[0][0], right = matrix[n - 1][n - 1];
        while (left < right) {
            int mid = (left + right) / 2;
            if (check(matrix, mid, k)) left = mid + 1;
            else right = mid;
        }
        return left;
    };
private:
    bool check(vector<vector<int>>& matrix, int cur, int k) {
        int cnt = 0, n = matrix.size(), i = n - 1, j = 0;
        while (i >= 0 && j < n) {
            if (matrix[i][j] <= cur) {
                j++;
                cnt += i + 1;
            }
            else i--;
        }
        return cnt < k;
    }
};

重点思路

由于本题只需要求解第k小的元素，而不是前k小的元素，所以可以不用像解答1一样枚举出前k小的元素，只需要二分答案即可。时间复杂度为\(O(nlog(max - min))\)。