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 to matrix). Now, the next smallest number is clearly matrix (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 (
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,  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