Kth Smallest Element in a Sorted Matrix LeetCode solution in Python

1
class Solution:
2
    def kthSmallest(self, matrix, k):
3
        n, beg, end = len(matrix), matrix[0][0], matrix[-1][-1]
4

5
        def check(m):
6
            i, j, cnt = 0, n - 1, 0
7
            for i in range(n):
8
                while j >= 0 and matrix[i][j] > m:
9
                    j -= 1
10
                cnt += j + 1
11
            return cnt
12

13
        while beg < end:
14
            mid = (beg + end) // 2
15
            if check(mid) < k:
16
                beg = mid + 1
17
            else:
18
                end = mid
19

20
        return beg