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
```