Kth Smallest Element in a Sorted Matrix — Binary Search on Value

# Python
def kthSmallest(matrix, k):
    n = len(matrix)
    def count_le(mid):
        # Count elements <= mid using staircase from bottom-left
        cnt = 0; r, c = n-1, 0
        while r >= 0 and c < n:
            if matrix[r][c] <= mid: cnt += r+1; c += 1
            else: r -= 1
        return cnt

    lo, hi = matrix[0][0], matrix[-1][-1]
    while lo < hi:
        mid = (lo+hi)//2
        if count_le(mid) >= k: hi = mid
        else: lo = mid+1
    return lo

// Java
public int kthSmallest(int[][] matrix, int k) {
    int n=matrix.length, lo=matrix[0][0], hi=matrix[n-1][n-1];
    while (lo<hi) {
        int mid=lo+(hi-lo)/2, cnt=0, r=n-1, c=0;
        while (r>=0&&c<n) { if(matrix[r][c]<=mid){cnt+=r+1;c++;} else r--; }
        if (cnt>=k) hi=mid; else lo=mid+1;
    }
    return lo;
}