Kth Smallest Element in a Sorted Matrix
Find the kth smallest value in a matrix where only rows and columns are individually sorted.
- n == matrix.length == matrix[i].length
- 1 <= n <= 300
- -10⁹ <= matrix[i][j] <= 10⁹
- All the rows and columns of matrix are guaranteed to be sorted in non-decreasing order.
- 1 <= k <= n²
Intuition
Kth smallest element in a sorted matrix works on a matrix whose rows and columns are each sorted ascending. Flattening and sorting is O(n² log n) and ignores the structure entirely.
Two better approaches exist, and they think about the problem differently.
The min-heap treats the matrix as n sorted lists being merged. Seed the heap with the first element of each row, then pop k times, pushing the next element from the popped element's row each time. The k-th pop is the answer, at O(k log n).
The binary search on value is the stronger method and scales better when k is large:
- Search the value range from matrix[0][0] to matrix[n-1][n-1], counting how many elements are less than or equal to each candidate.
That count is monotonic in the candidate, which is what permits binary search even though the candidate may not be an element of the matrix.
Counting exploits the sortedness in both directions. Start at the bottom-left corner: if the value there is at most the candidate, the entire column above it also qualifies, so add row + 1 and move right; otherwise move up. That staircase walk counts in O(n) rather than O(n²).
The search converges on the smallest value whose count reaches k, and that value is guaranteed to be in the matrix — a subtlety worth stating, since the midpoint being an arbitrary integer makes it look otherwise. The convergence point is the first value achieving the count, which must therefore be present.
Binary search gives O(n log(max − min)), independent of k.
A k-way merge across the rows. Seed the heap with each row's first element, then popping k times walks the merged order — each pop pushes only the next element from that same row, so the heap never exceeds n entries.
Approach
Before reading on: price up what the direct approach costs here, then ask why only the extreme element matters and the rest need no ordering. Aim for O(k log n) time and O(n) space.
Use the sorted structure
Flattening and sorting is O(n² log n) and discards what the problem provides. Both rows and columns are sorted, which two better approaches exploit.
Consider the min-heap merge
Seed a heap with each row's first element, then pop k times, pushing the next element from the popped row. O(k log n) — good when k is small.
Binary search the value range
Search from matrix[0][0] to matrix[n-1][n-1], counting elements at most each candidate. The count is monotonic, which is what makes this valid.
Count with a staircase walk
Start at the bottom-left corner. If the value is at most the candidate, add row + 1 and move right; otherwise move up. This counts in O(n), not O(n²).
Converge on the answer
When the count reaches k, search lower; otherwise higher. The loop settles on the smallest value whose count reaches k.
Know why the result is in the matrix
The converged value must be present, since it is the first value achieving the count — even though intermediate midpoints may not be matrix elements at all.
Cost of each approach
The heap is O(k log n); binary search is O(n log(max − min)), independent of k and preferable when k approaches n².
Solution & live demo
Common pitfalls
Pushing the whole matrix into the heap
heap = [v for row in matrix for v in row] heapq.heapify(heap)
heap = [(matrix[r][0], r, 0) for r in range(n)]
That's O(n²) space and discards the row ordering the problem gives you. Seeding one element per row keeps the heap at size n and pulls in the rest lazily.
Not tracking the column index
heapq.heappush(heap, matrix[r][c + 1])
heapq.heappush(heap, (matrix[r][c + 1], r, c + 1))
Without the coordinates in the entry, the next pop has no way to know which row to advance. The triple carries the position along with the value.
Popping k + 1 times
for _ in range(k + 1):
for _ in range(k):
The k-th smallest is the value of the k-th pop, since the first pop yields the smallest. One extra iteration returns the (k+1)-th element.
Edge cases
the single element is both the heap seed and the only possible answer for any valid k
the loop pops every element exactly once, ending on the matrix maximum
duplicates are treated as distinct heap entries and counted separately toward k
the heap never holds more than n entries at once, which is what keeps this cheaper than flattening and sorting everything