【Leetcode】解题笔记 1 （归并排序）——378.有序矩阵中第k小的元素

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$