LeetCode 378. 有序矩阵中第K小的元素

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

示例:

matrix = [
   [ 1,  5,  9],
   [10, 11, 13],
   [12, 13, 15]
],
k = 8,

返回 13。
说明:
你可以假设 k 的值永远是有效的, 1 ≤ k ≤ n^2 。

 

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int n = matrix.size();
        int low = matrix[0][0], high = matrix[n - 1][n - 1], mid, cnt;
        while (low < high)
        {
            mid = (low + high) / 2;
            cnt = countSmaller(matrix, mid);
            if (cnt < k)
                low = mid + 1;
            else
                high = mid;
        }
        return low;
    }

    int countSmaller(vector<vector<int>>& matrix, int key) {
        int n = matrix.size();
        int i = n - 1, j = 0;
        int cnt = 0;
        while (i >= 0 && j < n)
        {
            if (matrix[i][j] > key)
                --i;
            else
            {
                cnt += i + 1;
                ++j;
            }
        }
        return cnt;
    }
};