Search a 2D Matrix: Staircase Walk vs Binary Search

"Search a 2D matrix" is two different interview problems sharing one name. In one, the whole matrix is really a single sorted list folded into rows. In the other, only the rows and the columns are sorted, and binary search on the whole thing is not possible. The first question to ask is which one you have.

The problem it solves

Both versions give an m × n matrix of numbers and a target, and ask whether the target is present:

• Fully sorted. Each row is sorted, and each row's first value is greater than the previous row's last value. Read row by row, it is one sorted list.
• Row- and column-sorted. Each row increases left to right and each column increases top to bottom, but a row may start lower than the previous row ends. In [[1, 4, 7], [2, 5, 8]], the second row starts at 2, below 7.

The first allows a real binary search. The second does not, because there is no single order to bisect, and that is the case the lesson's animation is about.

The intuition

For the row- and column-sorted matrix, pick a starting cell where one comparison gives a clear instruction. The top-left corner does not: everything to its right and below it is larger, so "the target is bigger" points in two directions at once.

The top-right corner is different. It is the largest value in its row and the smallest in its column:

• If it is larger than the target, everything below it in that column is larger still, so the whole column is out. Step left.
• If it is smaller than the target, everything to its left in that row is smaller still, so the whole row is out. Step down.

Every comparison deletes a full row or a full column, so the walk ends within m + n steps. The bottom-left corner works too, with the directions mirrored.

For the fully sorted matrix, forget the grid. Index k of the flattened list lives at row k // n, column k % n, so an ordinary binary search over 0 .. m*n - 1 works without copying anything.

Watch it run

The animation searches a 3 × 3 matrix for 9. It starts on the top-right corner, 7. Seven is too small and the largest in its row, so the whole row goes. Next is 12, too big and the smallest in its column, so the column goes. Then 9: found in three steps instead of nine.

The code

The staircase walk for the row- and column-sorted matrix:

def staircase(matrix, target):
    if not matrix or not matrix[0]:
        return None
    r, c = 0, len(matrix[0]) - 1             # top-right: row maximum, column minimum
    while r < len(matrix) and c >= 0:
        v = matrix[r][c]
        if v == target:
            return (r, c)
        if v > target:
            c -= 1                           # everything below in this column is bigger
        else:
            r += 1                           # everything left in this row is smaller
    return None

m = [[1, 4, 7, 11],
     [2, 5, 8, 12],
     [3, 6, 9, 16],
     [10, 13, 14, 17]]
print(staircase(m, 9), staircase(m, 15))     # (2, 2) None

The flattened binary search for the fully sorted matrix:

def flat_binary_search(matrix, target):
    if not matrix or not matrix[0]:
        return None
    cols = len(matrix[0])
    lo, hi = 0, len(matrix) * cols - 1
    while lo <= hi:
        mid = (lo + hi) // 2
        v = matrix[mid // cols][mid % cols]  # flat index -> (row, column)
        if v == target:
            return divmod(mid, cols)
        if v < target:
            lo = mid + 1
        else:
            hi = mid - 1
    return None

g = [[1, 3, 5, 7], [10, 11, 16, 20], [23, 30, 34, 60]]
print(flat_binary_search(g, 16), flat_binary_search(g, 13))   # (1, 2) None

The same staircase idea counts how many cells are at most x, which is the building block for "k-th smallest in a sorted matrix". From the bottom-left, a cell that fits brings its whole column above it along:

def count_at_most(matrix, x):
    r, c, count = len(matrix) - 1, 0, 0
    while r >= 0 and c < len(matrix[0]):
        if matrix[r][c] <= x:
            count += r + 1                   # this cell and everything above it
            c += 1
        else:
            r -= 1
    return count

print(count_at_most(m, 9))                   # 9

All three against a plain scan, on 2,000 random matrices built so that every cell is at least its top and left neighbours, with repeated values allowed:

import random

def random_sorted_grid(rows, cols):
    g = [[0] * cols for _ in range(rows)]
    for r in range(rows):
        for c in range(cols):
            floor = max(g[r - 1][c] if r else 0, g[r][c - 1] if c else 0)
            g[r][c] = floor + random.randint(0, 3)
    return g

random.seed(15)
ok = True
for _ in range(2000):
    rows, cols = random.randint(1, 6), random.randint(1, 6)
    g = random_sorted_grid(rows, cols)
    t = random.randint(-1, max(map(max, g)) + 1)
    present = any(t in row for row in g)
    hit = staircase(g, t)
    ok &= (hit is not None) == present and (hit is None or g[hit[0]][hit[1]] == t)
    ok &= count_at_most(g, t) == sum(v <= t for row in g for v in row)
    flat = sorted(random.sample(range(100), rows * cols))
    fg = [flat[r * cols:(r + 1) * cols] for r in range(rows)]
    hit = flat_binary_search(fg, t)
    ok &= (hit is not None) == (t in flat) and (hit is None or fg[hit[0]][hit[1]] == t)
print(ok)                                    # True

The complexity

• Staircase: each step removes a row or a column, so at most m + n - 1 comparisons: O(m + n) time, O(1) space.
• Binary search per row on the row- and column-sorted matrix: O(m log n). Better than the staircase only when m is small relative to n.
• Flattened binary search on the fully sorted matrix: O(log(m * n)), which equals O(log m + log n). Using the staircase there still works, but throws away the stronger ordering.
• A full scan is O(m * n), the baseline the random check compares against.

Where it goes wrong

• Starting at the top-left or bottom-right. Both neighbours move in the same direction, so a comparison cannot rule anything out.
• Binary searching a matrix that is only row- and column-sorted. The flattened list is not sorted, so bisection skips the target. [[1, 4], [2, 5]] flattens to 1, 4, 2, 5.
• Mixing up the index maths. Row is k // cols and column is k % cols, using the number of columns, not rows.
• Empty input. [] and [[]] both need a guard before len(matrix[0]).

How to say it in an interview

"First I check which ordering I have. If every row starts above the previous row's end, it's one sorted list, and I binary search indexes 0 to m*n - 1, mapping k to row k // n, column k % n: O(log mn). If only rows and columns are sorted, I start at the top-right corner, the largest in its row and smallest in its column. Too big, I drop the column and move left; too small, I drop the row and move down. Each step removes a line, so it's O(m + n) with O(1) space."

The one-dimensional version is in binary search explained visually, and the counting walk above powers the k-th smallest in a sorted matrix.