Python程序在行排序矩阵中查找中位数

在一个行排序的矩阵中，找到其中位数。中位数是指将矩阵中所有元素按照大小排列，位于中间位置的元素。如果矩阵的元素个数是偶数，则中位数是矩阵中偏左的那个元素。

方法一：排序

最简单的方法就是将矩阵中的元素全部加入到一个列表中，然后排序后取中位数。

def kthSmallest(matrix, k):
    """
    :type matrix: List[List[int]]
    :type k: int
    :rtype: int
    """
    n = len(matrix)
    nums = []
    for i in range(n):
        for j in range(n):
            nums.append(matrix[i][j])
    return sorted(nums)[k-1]

时间复杂度为 $O(n^2logn)$，空间复杂度为 $O(n^2)$。

方法二：二分查找

二分查找是一种时间复杂度比较低的算法，我们可以用它来搜索矩阵中的中位数。因为矩阵已经排好序了，所以我们可以通过调整一个左右边界，来找到中间位置的元素。

def kthSmallest(matrix, k):
    """
    :type matrix: List[List[int]]
    :type k: int
    :rtype: int
    """
    m, n = len(matrix), len(matrix[0])

    def count_no_more_than(x):
        """
        统计矩阵中小于或等于 x 的元素个数
        """
        i, j = 0, n - 1
        cnt = 0
        while i < m and j >= 0:
            if matrix[i][j] <= x:
                cnt += j + 1
                i += 1
            else:
                j -= 1
        return cnt

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

时间复杂度为 $O(nlogn)$，空间复杂度为 $O(1)$。

总结

此题实现的方法还有很多，可以继续探索，但以上两种方法应该是较为常用的。这里仅作为一份介绍和参考。