Kth Smallest Element in a Sorted Matrix

A matrix whose rows and columns are each sorted ascending. Find the kth smallest value. Flattening and sorting would be O(n² log n²); the sorted structure exists to be exploited.

The key insight: the top-left cell is the global minimum, and each cell's right and down neighbors are its only possible successors in the sorted order. That makes the matrix a DAG where a min-heap can run a merge: seed the heap with the first cell of each row, pop the minimum, push the next cell in that row, and count. The kth pop is the answer.

Approach in steps:

1. Push (matrix[r][0], r, 0) for each row.
2. Pop the minimum k times; after each pop, push the next column cell in that row.
3. Return the value of the kth pop.

import heapq
 
def kthSmallest(matrix, k):
    n = len(matrix)
    heap = [(matrix[r][0], r, 0) for r in range(min(n, k))]
    heapq.heapify(heap)
    for _ in range(k):
        val, r, c = heapq.heappop(heap)
        if c + 1 < n:
            heapq.heappush(heap, (matrix[r][c + 1], r, c + 1))
    return val

Time is O(k log n) with an n×n matrix, space O(n). A binary-search-on-value alternative runs O(n log(max-min)) and suits huge k.

Trickiest edge case: duplicate values across cells. The heap stores (value, row, col) triples, so equal values stay distinct entries and the machine never collapses two matrix positions into one. Also guard the row-push when k < n — pushing all rows is fine but wasteful; seeding only min(n, k) rows is enough because no row beyond the k-th can contribute before the k-th pop.