Python heap based solution - 360ms


  • 0
    H
    class Solution(object):
        def kthSmallest(self, matrix, k):
            """
            :type matrix: List[List[int]]
            :type k: int
            :rtype: int
            """
            n = len(matrix)
            h = []  # heap
            r = 0   # rank
            seen = set()    # set of counted elements to avoid double entries in heap
            heapq.heappush(h, (matrix[0][0], 0, 0))
            while (r < k):
                num, i, j = heapq.heappop(h)
                # ignore if already counted
                if ((i, j) in seen):
                    continue
                seen.add((i, j))    # add to the seen set
                r += 1
                if (j < n - 1):
                    heapq.heappush(h, (matrix[i][j + 1], i, j + 1))
                if (i < n - 1):
                    heapq.heappush(h, (matrix[i + 1][j], i + 1, j))
                    
            return num
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.