Solution: Kth Smallest Number in a Sorted Matrix

from heapq import heappop, heappush

class Solution:
  def countSmallerEq(self, matrix, val):
    n, m = len(matrix), len(matrix[0]) # nxm
    
    row, col = 0, m - 1
    smaller = 0

    while row < n and col >= 0:
      if matrix[row][col] <= val:
        smaller += (col + 1)
        row += 1
      elif matrix[row][col] > val:
        col -= 1
    
    return smaller
    
  def findKthSmallest(self, matrix, k):

    n, m = len(matrix), len(matrix[0])
    lo = matrix[0][0]
    hi = matrix[n-1][m-1]

    while lo <= hi:
      mid = (lo + hi) // 2
      n_smaller = self.countSmallerEq(matrix, mid)
      if n_smaller >= k:
        hi = mid - 1
      else:
        lo = mid + 1
    
    return lo

For the counting step in the binary search method, I recommend looking at https://leetcode.com/problems/search-a-2d-matrix-ii/