LeetCode 378 有序矩阵中第K小的元素

给定一个 n x n 矩阵，其中每行和每列元素均按升序排序，找到矩阵中第 k 小的元素。
这道题和之前的搜索二维矩阵2类似。都是从矩阵的左下角开始遍历。第k小的元素有一个范围，在最大值和最小值之间进行二分搜索。

from typing import *

class Solution:

    def check(self, matrix, val):
        '''
        找到矩阵中小于等于val元素的个数
        :param matrix:
        :param val:
        :return:
        '''
        m, n = len(matrix), len(matrix[0])
        i, j = m - 1, 0
        cnt = 0
        while i >= 0 and j < n:
            if matrix[i][j] <= val:
                cnt += (i + 1)
                j += 1
            else:
                i -= 1
        return cnt

    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        l, r = matrix[0][0], matrix[-1][-1]
        while l <= r:
            mid = (l + r) // 2
            num = self.check(matrix, mid)
            if num >= k:  # 找到的为最小的l，最小的l正好在数组中
                # 大于等于时要向左搜索
                r = mid - 1
            else:
                l = mid + 1
        return l