leetcode 有序矩阵中第K小的元素 python

import heapq
class Solution(object):
    def kthSmallest(self, matrix, k):
        nums=heapq.merge(*matrix)
        for i in range(k-1):
            nums.next()
        return nums.next()

别问我，我也不知道发生了什么
使用堆“优化”了一番

import heapq
class Solution(object):
    def kthSmallest(self, matrix, k):
        row_num=len(matrix)
        col_num=len(matrix[0])
        heap=[]
        heapq.heapify(heap)
        def heap_add(num):
            if len(heap)>=k:
                heapq.heappop(heap)
            heapq.heappush(heap,-num)
        for row in matrix:
            if len(heap)>=k and row[0]>=-heap[0]:
                break
            for col in row:
                if len(heap)>=k and col>=-heap[0]:
                    break
                heap_add(col)
        return -heap[0]

速度还没之前快，真他妈不科学啊