力扣—378有序矩阵第k小的数

题目描述

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

示例：

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

输出：13

解题思路

由于每行均为递增序列，而且是n x n的方阵，所以可以维护一个小根堆，首先取第一列数据[1, 10, 12]，维护一个小根堆（堆顶元素为1），然后从堆弹出k - 1次，那么下一次弹出的就是第k小的元素，假设弹出1后，判断是否到了最后一列，如果没有到，那么往堆中加入同行下一个元素，这里是加入5，所以新的堆为[5, 10, 12]，再弹出5，以此类推。

import heapq
class Solution(object):
    def kthSmallest(self, matrix, k):
        n = len(matrix)
        pq = [(matrix[i][0], i, 0) for i in range(n)]
        heapq.heapify(pq)
        ret = 0
        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]

S = Solution()
print(S.kthSmallest([[ 1,  5,  9],
                     [10, 11, 13],
                     [12, 13, 15]], 8))