二分法 有序矩阵中第K小的元素

leetcode 378
题目：
给定一个 n x n 矩阵，其中每行和每列元素均按升序排序，找到矩阵中第 k 小的元素。
请注意，它是排序后的第 k 小元素，而不是第 k 个不同的元素。
matrix = [
[ 1, 5, 9],
[10, 11, 13],
[12, 13, 15]
],
k = 8,
返回 13。

def kthSmallest( matrix ,k):
    n = len(matrix)
    def search(mid):  # 记录矩阵中小于等于mid的个数
        i = n - 1  # i初始为最后一行，左下角的元素
        j = 0  # j初始为第一列
        count = 0
        # 以列为单位找，找到每一列最后一个<=mid的数即知道每一列有多少个数<=mid
        while j < n and i >= 0:
            if matrix[i][j] <= mid:
                count += i + 1
                j += 1
            else:
                i -= 1
        return count

    left = matrix[0][0]
    right = matrix[-1][-1]
    while left < right:
        mid = (left + right) >> 1
        num = search(mid)
        if num < k:
            left = mid + 1
        else:
            right = mid
    return left