Kth Smallest Element in a Sorted Matrix

lcpy gen -n 378   # by problem number
lcpy gen -s kth_smallest_element_in_a_sorted_matrix   # by problem name

Input: matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
Output: 13
Explanation: The elements in the matrix are [1,5,9,10,11,12,13,13,15], and the 8th smallest number is 13

Input: matrix = [[-5]], k = 1
Output: -5

class Solution:
    # Time: O(n * log(max - min))
    # Space: O(1)
    def kth_smallest(self, matrix: list[list[int]], k: int) -> int:
        def count_at_most(target: int) -> int:
            count = 0
            row, col = len(matrix) - 1, 0
            while row >= 0 and col < len(matrix):
                if matrix[row][col] <= target:
                    count += row + 1
                    col += 1
                else:
                    row -= 1
            return count

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