用堆来解决
调用标准库import heapq
先把矩阵最左端一列压入队中
初始化结果
接着循环k次,每次循环把堆顶元素弹出,再压入弹出元素右边的元素(如果存在的话)
class Solution(object):
def kthSmallest(self, matrix, k):
"""
:type matrix: List[List[int]]
:type k: int
:rtype: int
"""
heap = [(row[0], i, 0) for i, row in enumerate(matrix)]
heapq.heapify(heap)
ret = 0
for _ in range(k):
ret, i, j = heapq.heappop(heap)
if j+1 < len(matrix[0]):
heapq.heappush(heap, (matrix[i][j+1], i, j+1))
return ret