8 lines simple O(min(K, N-K) logN) solution with comments

  • 0

    Similar idea to other O(KlogK) solution, but I only put possible next smallest number into the heap. A simple example to illustrate the difference: assume N=4, and right now I've found 4 smallest numbers, and they are just the 4 numbers of the first row (matrix[0][1] to matrix[0][3]). Now, the next smallest number is clearly matrix[1][0] (only possibility), but not all the numbers adjacent to those 4 numbers. Now, the numbers in the heap is obviously upper-bounded by N, and thus O(KlogN) instead of O(KlogK).

    To realize this, I log for each row and each column, what is the maximum index I've already used (r and c).

    Another optimization we can do is that if K is big, then we can do N-K instead, by going from the biggest number (thus, min(K, N-K) instead of K). I didn't do this in my code, which should be easy to implement though.

        def kthSmallest(self, matrix, k):
            # r, c log for smallest index I've used for each row and column
            r, c = [-1 for x in range(len(matrix))], [-1 for x in range(len(matrix))] 
            # x, y is current coordinates, v is current (smallest) value, h is just the heap
            x, y, v, h = 0, 0, matrix[0][0], []
            for _ in range(k - 1):
                # update r, c
                r[x], c[y] = y, x 
                # check whether I should put [x+1][y] to the heap
                if x + 1 < len(matrix) and r[x + 1] == y - 1: heapq.heappush(h, (matrix[x + 1][y], (x + 1, y))) 
                if y + 1 < len(matrix) and c[y + 1] == x - 1: heapq.heappush(h, (matrix[x][y + 1], (x, y + 1)))
                v, (x, y) = heapq.heappop(h)
            return v

Log in to reply

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