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

给你一个

n x n

矩阵

matrix

，其中每行和每列元素均按升序排序，找到矩阵中第

k

小的元素。

请注意，它是 排序后 的第

k

小元素，而不是第

k

个 不同 的元素。

输入：matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
输出：13
解释：矩阵中的元素为 [1,5,9,10,11,12,13,13,15]，第 8 小元素是 13
           

最小堆

由题目给出的性质可知，这个矩阵的每一行均为一个有序数组。问题即转化为从这 n n n 个有序数组中找第 k k k​ 大的数.

首先把第一列构造最小堆，然后依次删除堆中最小值

num

，并加入矩阵剩余的最小值。剩余的最小值是

num

右移一格的值，最小堆（蓝色的值）的变换如下：

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix) #注：题目中这个矩阵是n*n的，所以长宽都是n

        col = [(matrix[i][0], i ,0) for i in range(n)] #取出第一列的值
        # matrix[i][0]是具体的值，后面的(i,0)是在记录矩阵中的位置，方便每次右移添加下一个
        heapq.heapify(col) #变成一个heap
        for i in range(k - 1): #一共弹k次：这里k-1次，return的时候1次
            num, x, y = heapq.heappop(col) #弹出堆里最小一个
            if y != n - 1: #如果这一行还没被弹完
                heapq.heappush(col , (matrix[x][y+1], x ,y+1 )) #加入num右边的一个值
        return heapq.heappop(col)[0]


           

二分法

1. 找出二维矩阵中最小的数 l e f t l e f t left​, 最大的数 right, 那么第 k k k​ 小的数必定在 left ∼ \sim ∼​ right 之间.
2. mid = ( =( =(​​ left + + +​​ right ) / / 2 ) // 2 )//2​​; 在二维矩阵中寻找小于等于 m i d m i d mid​​ 的元素个数 count.
3. 若这个 count 小于 k k k, 表明第 k k k 小的数在右半部分且不包含 mid ⁡ \operatorname{mid} mid, 即 l e f t = m i d + 1 , r i g h t = l e f t=m i d+1, r i g h t= left=mid+1,right= right, 保证了第 k k k 小的数在 l e f t ∼ r i g h t l e f t \sim right left∼right 之间
4. 若这个 count >= k k k, 表明第 k k k 小的数在左半部分且可能包含 mid, 即 l e f t = l e f t , r i g h t = l e f t=l e f t, r i g h t= left=left,right= mid ⁡ \operatorname{mid} mid, 又保证了第 k k k 小的数在 l e f t ∼ r i g h t l e f t \sim right left∼right之间
5. 因为每次循环中都保证了第 k k k 小的数在 l e f t ∼ r i g h t l e f t \sim right left∼right之间，当 l e f t = = r i g h t l e f t==r i g h t left==right 时，第 k k k 小的数 即被找出，等于 r i g h t \mathrm{right} right

注意：这里的 l e f t 、 m i d 、 r i g h t left 、mid 、right left、mid、right 是数值，不是索引位置。

count 的计算如下，取 m i d = 8 mid =8 mid=8.

可以这样描述走法：

• 初始位置在 m a t r i x [ n − 1 ] [ 0 ] matrix[n - 1][0] matrix[n−1][0]​（即左下角）；
• 设当前位置为 matrix ⁡ [ i ] [ j ] \operatorname{matrix}[i][j] matrix[i][j] 。若 matrix ⁡ [ i ] [ j ] ≤ m i d \operatorname{matrix}[i][j] \leq m i d matrix[i][j]≤mid, 则将当前所在列的不大于 m i d m i d mid 的数的数量 (即 i + 1 i+1 i+1 ) 累加到

  count

  中 , 并向右移动, 否则向上移动;
• 不断移动直到走出格子为止。

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        row = len(matrix)
        col = len(matrix[0])
        left = matrix[0][0]
        right = matrix[row - 1][col - 1]
        while left < right:
            #每次循环都保证第K小的数在left~right之间，当left==right，第k小的数就是left
            mid = left+(right-left)//2
            count = self.findNotBiggerThanMid(matrix, mid, row, col)
            if count<k:
                #第k小的数在右半部分，且不包含mid
                left = mid + 1
            else :
                #第k小的数在左半部分，可能包含mid
                right = mid
        return right

    def findNotBiggerThanMid(self,matrix,mid,row,col):
        #以列为单位找，找到每一列最后一个<=mid的数即知道每一列有多少个数<=mid
        i, j = row - 1, 0
        count =0
        while i >= 0 and j < col:
            if matrix[i][j] <= mid:
                count += i + 1
                j += 1
            else:
                i -= 1
        return count
           

参考

Krahets - 力扣（LeetCode） (leetcode-cn.com)