获取第k个元素

问题

n*n矩阵按行、按列升序排列，获取第k个小的元素。

方法

1.二分查找

由题目给出的性质可知，这个矩阵内的元素是从左上到右下递增的（假设矩阵左上角为 matrix[0][0]matrix[0][0]matrix[0][0]）。

我们知道整个二维数组中 matrix[0][0]matrix[0][0]matrix[0][0] 为最小值，matrix[n−1][n−1]matrix[n - 1][n - 1]matrix[n−1][n−1] 为最大值，现在我们将其分别记作 l 和 r。

可以发现一个性质：任取一个数 mid 满足 l≤mid≤r ，那么矩阵中不大于 mid 的数，肯定全部分布在矩阵的左上角。我们只要沿着左上角的这条锯齿线走一遍即可计算出这两个板块的大小，也自然就统计出了这个矩阵中不大于 midmidmid 的数的个数了

可以这样描述走法：

初始位置在 matrix[n−1][0]matrix[n - 1][0]matrix[n−1][0]（即左下角）；

设当前位置为 matrix[i][j]matrix[i][j]matrix[i][j]。若 matrix[i][j]≤midmatrix[i][j] \leq midmatrix[i][j]≤mid，则将当前所在列的不大于 midmidmid 的数的数量（即 i+1i + 1i+1）累加到答案中，并向右移动，否则向上移动；

不断移动直到走出格子为止。

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)

        def check(mid):
            i, j = n - 1, 0
            num = 0
            while i >= 0 and j < n:
                if matrix[i][j] <= mid:
                    num += i + 1
                    j += 1
                else:
                    i -= 1
            return num >= k

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

2.最小堆-归并排序

由题目给出的性质可知，这个矩阵的每一行均为一个有序数组。问题即转化为从这 nnn 个有序数组中找第 kkk 大的数，可以想到利用归并排序的做法，归并到第 kkk 个数即可停止。

一般归并排序是两个数组归并，而本题是 nnn 个数组归并，所以需要用小根堆维护，以优化时间复杂度。

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)
        pq = [(matrix[i][0], i, 0) for i in range(n)]
        heapq.heapify(pq)

        ret = 0
        for i in range(k - 1):
            num, x, y = heapq.heappop(pq)
            if y != n - 1:
                heapq.heappush(pq, (matrix[x][y + 1], x, y + 1))
        
        return heapq.heappop(pq)[0]

【1】. - 力扣（LeetCode）