378 - binary search

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)
        l = matrix[0][0] - 1
        r = matrix[-1][-1]
        res = r

        def count_less_than_x(x):
            count = 0
            max_elem = 0
            i, j = n - 1, 0
            while i >= 0 and j < n:
                if matrix[i][j] > x:
                    i -= 1
                else:
                    count += (i + 1)
                    max_elem = max(max_elem, matrix[i][j])
                    j += 1
            return count, max_elem


        while r - l > 1:
            mid = l + (r - l) // 2
            count, val = count_less_than_x(mid)
            if count < k:
                l = mid
            else:
                r = mid
                res = val

        return res