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

Posted 2021-09-28 炫云云

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

注意：这里的 $left、mid、right$ 是数值，不是索引位置。

count 的计算如下，取$mid=8$.

可以这样描述走法：

• 初始位置在 $matrix[n-1][0]$​（即左下角）；
• 设当前位置为 $\mathrm{matrix}[i][j]$ 。若 $\mathrm{matrix}[i][j]\le mid$, 则将当前所在列的不大于 $mid$ 的数的数量 (即 $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)