Find Kth Smallest Element in a Sorted Matrix

Intuition

Instead of sorting or scanning entries directly, guess a value and ask how many grid entries do not exceed it. Because both rows and columns are sorted, that count can be gathered in linear time with a staircase walk that starts at the bottom-left corner: moving right adds a whole column prefix, moving up skips values that are too large. The count of entries at or below a guess never decreases as the guess grows, so we binary search for the smallest value whose count reaches k. That value is guaranteed to be an actual grid entry.

Algorithm walkthrough

1. 1Set lo to the top-left entry (the global minimum) and hi to the bottom-right entry (the global maximum).
2. 2Pick mid between lo and hi and count entries less than or equal to mid.
3. 3Do the counting with a pointer at the bottom-left: if the current entry is at most mid, add the whole column above it and move right; otherwise move up.
4. 4If the count is at least k, the answer is at most mid, so set hi = mid; otherwise set lo = mid + 1.
5. 5When lo equals hi, that value is the k-th smallest element.

Implementation

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)

        def countLessEqual(target: int) -> int:
            count = 0
            row, col = n - 1, 0
            while row >= 0 and col < n:
                if matrix[row][col] <= target:
                    count += row + 1
                    col += 1
                else:
                    row -= 1
            return count

        lo, hi = matrix[0][0], matrix[n - 1][n - 1]
        while lo < hi:
            mid = (lo + hi) // 2
            if countLessEqual(mid) >= k:
                hi = mid
            else:
                lo = mid + 1
        return lo