Kth Smallest Element in a Sorted Matrix: Binary Search on Values

You get an n × n matrix whose rows and columns are both sorted, and a number k. Return the k-th smallest value. Sorting all n² values works and ignores the structure. A heap merging the rows works and is the answer most people find first. The answer that impresses is stranger: binary search over values rather than positions, with a counting step that walks the matrix like a staircase. It is also the cleanest example of binary search on the answer you will meet.

The problem it solves

The matrix [[1, 5, 9], [10, 11, 13], [12, 13, 15]] has sorted rows and sorted columns, but reading it row by row does not give a sorted list: 12 in the last row is smaller than 13 in the middle row. So there is no index to binary search. The 8th smallest value is 13, counting duplicates, and the task is to find it without sorting all n² values.

Three approaches, in increasing order of cleverness:

• Flatten and sort: O(n² log n) time and O(n²) space.
• Min-heap of row heads: push the first value of each row, pop k - 1 times, pushing the next value in the popped row each time: O(k log n).
• Binary search on values: O(n log(max - min)), with O(1) extra space.

The intuition

Instead of asking "which position holds the answer?", ask "how many values are at most x?" That count is easy to compute and it only grows as x grows. The answer is the smallest x for which the count reaches k.

The value range runs from the top-left corner, the smallest value, to the bottom-right corner, the largest. Guess the middle of the range, mid, and count the entries at or below it:

• Count below k: too few values are that small, so the answer is higher. Set lo = mid + 1.
• Count at least k: the answer is mid or lower, and mid itself may be the answer. Set hi = mid.

Counting uses the same staircase walk as searching a sorted 2D matrix. Start at the top-right corner. In each row, move the column pointer left while the value is above mid; everything left of the pointer is at most mid, so add column + 1 to the count. Because columns are sorted, the pointer never needs to move right again when you step down a row, so the whole count is O(n).

Why is the final value really in the matrix, when the search ranged over arbitrary integers? The loop ends at the smallest x whose count reaches k. If x were not in the matrix, then x - 1 would have the same count, contradicting "smallest". The same template, a monotone yes/no question over a range of candidate answers, is the one on the patterns cheat sheet under binary search.

Watch it run

The animation opens with the problem: the rows are sorted and so are the columns, but there is no index to binary search, so it searches the values instead. The range is 1 to 15. Guess 8: only 2 entries are at or below it, and 8 are needed, so the answer is higher. Guess 12: only 6 entries are at or below it, still too few, so the answer is higher again. Guess 14: 8 entries sit at or below it, enough, so the answer is this value or lower. Guess 13: 8 entries again, enough, so the answer is 13 or lower. The range closes on 13, and the final frame makes the key point: a value the counts land on is always a value that is really in the matrix.

The code

The lesson's search with the count pulled out into a function and a trace of each guess. The trace matches the animation's four guesses:

def count_at_most(matrix, x):
    """Staircase walk from the top-right corner: O(n) for an n x n matrix."""
    n = len(matrix)
    count, c = 0, n - 1
    for r in range(n):
        while c >= 0 and matrix[r][c] > x:
            c -= 1                          # the column cursor only ever moves left
        count += c + 1                      # everything left of c in this row is <= x
    return count

def kth_smallest(matrix, k, trace=None):
    n = len(matrix)
    lo, hi = matrix[0][0], matrix[n - 1][n - 1]
    while lo < hi:
        mid = (lo + hi) // 2
        count = count_at_most(matrix, mid)
        if trace is not None:
            trace.append((mid, count))
        if count < k:
            lo = mid + 1                    # too few at or below mid: answer is higher
        else:
            hi = mid                        # enough: mid itself may be the answer
    return lo

M = [[1, 5, 9], [10, 11, 13], [12, 13, 15]]
guesses = []
print(kth_smallest(M, 8, guesses))      # 13
print(guesses)                          # [(8, 2), (12, 6), (14, 8), (13, 8)]
print([kth_smallest(M, k) for k in range(1, 10)])   # [1, 5, 9, 10, 11, 12, 13, 13, 15]

The heap alternative, which merges the rows like a k-way merge, plus two edge cases for the value search: negative values, where Python's floor division still rounds mid down, and a matrix of duplicates:

import heapq

def kth_smallest_heap(matrix, k):
    n = len(matrix)
    heap = [(matrix[r][0], r, 0) for r in range(min(n, k))]   # the head of each row
    heapq.heapify(heap)
    for _ in range(k - 1):
        _, r, c = heapq.heappop(heap)
        if c + 1 < n:
            heapq.heappush(heap, (matrix[r][c + 1], r, c + 1))
    return heap[0][0]

print(kth_smallest_heap(M, 8))          # 13
print(kth_smallest([[-5, -4], [-5, 3]], 2), kth_smallest([[2, 2], [2, 2]], 3))   # -5 2

Both versions checked against the brute force, flattening and sorting, on 2,000 seeded random matrices up to 8 × 8 with duplicates and negatives, and the staircase count checked against a direct count:

import random

def random_sorted_matrix(n):
    m = [[0] * n for _ in range(n)]
    for r in range(n):
        for c in range(n):
            up = m[r - 1][c] if r else -30
            left = m[r][c - 1] if c else -30
            m[r][c] = max(up, left) + random.randint(0, 4)   # rows and columns never decrease
    return m

random.seed(23)
ok = True
for _ in range(2000):
    n = random.randint(1, 8)
    m = random_sorted_matrix(n)
    flat = sorted(v for row in m for v in row)                # brute force: sort everything
    k = random.randint(1, n * n)
    ok &= kth_smallest(m, k) == kth_smallest_heap(m, k) == flat[k - 1]
    x = random.randint(-35, 60)
    ok &= count_at_most(m, x) == sum(v <= x for v in flat)
print(ok)                               # True

The complexity

With an n × n matrix and a value range R = max - min:

• Binary search on values: O(n log R) time, O(1) extra space. Each of the log R rounds costs one O(n) staircase count.
• Heap: O(k log n) time and O(min(n, k)) space, since only the first k rows can matter. Better when k is small.
• Flatten and sort: O(n² log n) time and O(n²) space.

Where it goes wrong

• Setting hi = mid - 1 when the count reaches k. mid may be the answer, and this discards it.
• Counting < mid instead of <= mid. The search then lands on the answer plus one, which may not even be in the matrix.
• Binary searching indices. The flattened matrix is not sorted, so an index search has nothing to halve.
• Using (lo + hi) / 2. True division produces floats, and the search stops landing on integers from the matrix.
• Floating-point values. The value search relies on integer steps; with real numbers, search until the range is tiny or use the heap.

When it shows up in interviews

It is a standard medium question on matrices and binary search, and a common follow-up to searching a sorted 2D matrix. The value-search idea is the one from binary search on the answer, and the same template solves "k-th smallest pair distance" and "k-th smallest number in a multiplication table". The heap version belongs with the top-k problems.

How to say it in an interview

"The matrix is not globally sorted, so I cannot binary search positions, but the number of entries at most x grows with x, so I binary search values between the top-left and bottom-right corners. For each guess I count with a staircase walk from the top-right corner, which is O(n) because the column pointer only moves left. Fewer than k means the answer is higher; k or more means it is the guess or lower, so I keep the guess. The search ends on the smallest value with a count of at least k, which must be in the matrix. That is O(n log R) time and O(1) space; for small k a heap of row heads is O(k log n)."