Leetcode 378 - Kth Smallest Element in a Sorted Matrix（二分+杨氏矩阵）

题意

求杨氏矩阵内的第k大元素。

思路

我们首先二分一下第k大元素，假设为x。然后只需要求x在杨氏矩阵内的排名就好（有多少个元素比x小）。

求元素x在杨氏矩阵内的排名，时间复杂度$O(m + n)$：

1. 初始order = 0, 从右上角$a[i, j]$开始。
2. 若$x \geq a[i, j]: i++, order += j + 1$（加上这一行所有$\leq x$的个数）
3. 若$x < a[i, j]: j--$

时间复杂度：$O((m + n)(log(a[m - 1, n - 1] - a[0, 0])))$

代码

class Solution {
public:
    vector<vector<int>> a;
    int rank(int x) {
        int order = 0, i = 0, j = a[0].size() - 1;
        while (i < a.size() && j >= 0) {
            if (x >= a[i][j]) order += j + 1, i++;
            else j--;
        }
        return order;
    }

    int kthSmallest(vector<vector<int>>& matrix, int k) {
        a = matrix;
        int m = a.size(), n = a[0].size();
        int L = a[0][0], R = a[m - 1][n - 1], M;
        while (L < R) {
            M = L + (R - L >> 1);
            if (rank(M) < k) L = M + 1;
            else R = M;
        }
        return R;
    }
};