1272 · 有序矩阵中的第K小元素

1272 · 有序矩阵中的第K小元素

int n;

public int kthSmallest(int[][] matrix, int k) {
    this.n = matrix.length;
    int lo = matrix[0][0], hi = matrix[n - 1][n - 1];
    while (lo < hi) {
        int mid = lo + hi >> 1;
        int count = check(matrix, mid);
        if (count < k) {
            lo = mid + 1;
        } else {
            hi = mid;
        }
    }
    return lo;
}


/**
 * 判断matrix中小于等于target的数量
 *
 * @param matrix
 * @param target
 * @return
 */
private int check(int[][] matrix, int target) {
    int ans = 0;
    int i = n - 1, j = 0;
    while (i >= -0 && j < n) {
        if (matrix[i][j] <= target) {
            ans += i + 1;
            j++;
        } else {
            i--;
        }
    }
    return ans;
}