378_KthSmallestElementInSortedMatrix.py 文件源码

def kthSmallest(self, matrix, k):
        """
        Use heap
        Time complexity: O(klog(n)+nlog(n))

        :type matrix: List[List[int]]
        :type k: int
        :rtype: int
        """        
        import heapq
        # (value, y, x)
        # value means item in the matrix
        # y means the y coordinate of the item
        # x means the x coordinate of the item
        h =  [(matrix[0][x], 0, x) for x in range(len(matrix[0]))] 
        heapq.heapify(h)
        for i in range(k):
            val, y, x = heapq.heappop(h)
            if y + 1 < len(matrix):
                heapq.heappush(h, (matrix[y+1][x], y+1, x))

        return val