使用小顶堆找出有序矩阵中第 K 小的元素

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

import heapq

def kthSmallest(matrix, k):
    heap = []
    n = len(matrix)

    for i in range(n):  # 使用第一列数据构建小顶堆
        heapq.heappush(heap, (matrix[i][0], i, 0))

    # 弹出k-1次
    k1 = 1
    while k1 < k:
        num, i, j = heapq.heappop(heap)
        if j < n - 1:  # 把第i行的所有数据加入小顶堆
            heapq.heappush(heap, (matrix[i][j + 1], i, j + 1))
        k1 += 1

    # 弹出第k次，也就得到第k小的数
    ans, _, _ = heapq.heappop(heap)
    return ans


matrix = [[1, 5, 9], [10, 11, 13], [12, 13, 15]]
k = 8
ans = kthSmallest(matrix, k)
print(ans) # 13