【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 ≤ k ≤ 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$