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

题目描述

给定一个 $n\times n$矩阵，其中每行和每列元素均按升序排序，找到矩阵中第$k$ 小的元素。 请注意，它是排序后的第 $k$ 小元素，而不是第$k$个不同的元素 示例

matrix = [ [ 1, 5, 9], [10, 11, 13], [12, 13, 15] ], k = 8, 返回 13

提示 你可以假设 $k$ 的值永远是有效的，$1\le k\le n^2$

解题思路——归并排序

矩阵中，每行list都是有序的，等价于$n$个list归并排序。因此根据归并排序思想，先获取矩阵第一列进行堆排序，提取出堆中最小的数和该数的行、列数，将该行的下一列数入堆排序，然后继续提取最小数，依次获取到$k$个最小值即为目标结果。

代码实现

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) 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]

复杂度分析

时间复杂度为$O(k\log n)$，归并k次，每次都进行插入和弹出，其时间消耗为$O(\log n)$ 空间复杂度为$O(n)$，堆排序消耗空间始终都为$n$ 注意 在最坏情况下$k$有可能为$n^2$