矩阵第K小

# -*- coding: utf-8 -*-


#
# 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
#
#
# @param matrix int整型二维数组
# @param k int整型
# @return int整型
#
class Solution:
    """
    题目：
        https://www.nowcoder.com/practice/c754e7a920614cba9b8b692ba9b20b5d?tpId=196&tqId=40447&rp=1&ru=/exam/oj&qru=/exam/oj&sourceUrl=%2Fexam%2Foj%3FjudgeStatus%3D3%26page%3D1%26pageSize%3D50%26search%3D%26tab%3D%25E7%25AE%2597%25E6%25B3%2595%25E7%25AF%2587%26topicId%3D196&difficulty=undefined&judgeStatus=3&tags=&title=
    参考：
        大神：牛客278599503号
    算法：
        题目已知：矩阵行列均为n，每行、每列均按升序排序，寻找第k小的元素。我们知道最小元素为left = matrix[0][0]，最大元素为right = matrix[-1][-1]，我们在闭合区间[left, right]上，使用二分法，寻找满足条件的第k小元素。
        函数check(target)：用于统计小于等于target的元素个数count（根据二维矩阵的升序性质，我们可以选择左下角或者右上角进行作为起点），返回值：count是否大于等于k
        当left < right时：
            取mid = left + (right - left) >> 1
            check(mid)：
                如果满足条件，说明第k小元素在区间[left, mid]中
                否则：说明mid小了，第k小元素在区间[mid + 1, right]中
    复杂度：
        时间复杂度：O(logD * (n + n)) = O(nlogD), D为矩阵最大值与最小值的差值，check方法的复杂度为O(2n) = O(n)
        空间复杂度：O(1)
    """

    def KthinMatrix(self, matrix, k):
        # write code here
        def check(target):
            i, j = 0, n - 1 # 这里选择从右上角开始统计
            count = 0

            while i < n and j >= 0:
                if matrix[i][j] <= target:
                    count += j + 1
                    i += 1
                else:
                    j -= 1

            return count >= k

        n = len(matrix)

        left, right = matrix[0][0], matrix[-1][-1]
        while left < right:
            mid = left + (right - left) / 2
            if check(mid):
                right = mid
            else:
                left = mid + 1
        return left


if __name__ == "__main__":
    sol = Solution()

    matrix, k = [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10],
                 [11, 12, 13, 14, 15], [16, 17, 18, 19, 20],
                 [21, 22, 23, 24, 25]], 5

    # matrix, k = [[1, 1, 1], [1, 1, 1], [1, 1, 1]], 5

    # matrix, k = [[1, 2, 5], [3, 4, 6], [7, 8, 9]], 5

    res = sol.KthinMatrix(matrix, k)

    print res