Searching 2D Matrices

A sorted matrix is usually a 1D pattern in disguise. First work out how it's sorted. If each row continues where the previous one ended, the matrix is one long sorted array. If rows and columns are sorted separately, use a corner where one move shrinks a whole row or column. For "Kth smallest" questions, binary search the values instead of positions.

The idea#

Most problems fall into one of three shapes, and each has its own technique.

Matrix shape	Technique	Cost (m rows, n cols)
Each row sorted, and each row starts after the previous row ends	treat it as one sorted array of length m·n	$O(\log(mn))$
Rows sorted and columns sorted, independently	staircase search from the top-right corner	$O(m + n)$
Rows and columns sorted, and you need the Kth smallest	binary search on the value, count cells ≤ guess	$O((m + n) \log(\text{range}))$

Flattening#

The first shape is the easiest, because the matrix already is one sorted array. Index i of the imaginary flat array lives at matrix[i // n][i % n]. Run the ordinary search on 0 .. m*n - 1. Searching the rows first (by each row's first and last value) and then the chosen row is the same idea in two steps.

Staircase search#

When rows and columns are sorted separately, the flat array isn't sorted, so flattening breaks. Instead, start at the top-right cell. Everything to its left in the row is smaller and everything below it in the column is larger, so one comparison throws away a whole line:

Cell bigger than the target: nothing below it in this column can match, so move left.
Cell smaller than the target: nothing to its left in this row can match, so move down.

Each step removes a row or a column, so it takes at most $m + n$ steps: $O(m + n)$ , not logarithmic. Binary searching every row instead costs $O(m \log n)$ .

How to spot it#

These cues suggest a matrix search.

"Search a 2D matrix", "each row is sorted", "each column is sorted".
Read the sorting rule closely: "the first integer of each row is greater than the last of the previous row" means flattening works.
"Kth smallest" in a sorted matrix or a multiplication table: search the value range.

Template#

Here are the two position-based searches. The value search for Kth smallest is covered in the hints. First the flattened search, for rows that chain.

m, n = number of rows, number of columns
l, r = 0, m * n - 1
while l <= r:
    mid = (l + r) // 2
    val = matrix[mid // n][mid % n]    # 1D index -> (row, col)
    if val == target: found
    elif val < target: l = mid + 1
    else: r = mid - 1
not found

Then the staircase, for rows and columns sorted separately.

row, col = 0, last column              # top-right corner
while row is in range and col >= 0:
    if matrix[row][col] > target: col -= 1   # rest of column is bigger
    elif matrix[row][col] < target: row += 1 # rest of row is smaller
    else: found
not found

The table traces the staircase for target 5 in a matrix whose first two rows are [1, 4, 7, 11, 15] and [2, 5, 8, 12, 19].

Step	Cell	Value vs target 5	Move
1	(0, 4)	15 > 5	left
2	(0, 3)	11 > 5	left
3	(0, 2)	7 > 5	left
4	(0, 1)	4 < 5	down
5	(1, 1)	5 = 5	found

Pitfalls#

Most mistakes come from misreading how the matrix is sorted.

Calling staircase search logarithmic. It walks at most m + n cells: $O(m + n)$ .
Starting at the wrong corner. Top-left and bottom-right give no decision, because both neighbors are bigger (or both smaller). Use top-right or bottom-left.
Flattening a matrix that doesn't chain. It only works when each row starts after the previous row ends.
Redundant bounds checks. With l <= r on row indices, the row index is always valid; extra min/max guards only hide bugs.

Problems in this pattern#

These problems use this pattern. Solve them before you open the hints below.

Hints#

Open a hint only after you have tried the problem for a while.

Hint: Search a 2D Matrix

Each row starts after the previous one ends, so the matrix chains into one sorted array. Flatten it and map i to (i // n, i % n).

Hint: Search a 2D Matrix II

Rows and columns are sorted separately, so flattening fails. Walk the staircase from the top-right corner.

Hint: Kth Smallest Element in a Sorted Matrix

Binary search the value, not the position. Count the cells ≤ a guess with the same staircase walk, and find the first value whose count is ≥ k. The "ghost guess" trap is covered in Binary Search on the Answer. A heap-based alternative is the K-way merge.

Hint: Kth Smallest Number in Multiplication Table

This is the same value search, with a formula for the count instead of a staircase walk. See Binary Search on the Answer.

Hint: Find a Peak Element II

Binary search over rows. Take the maximum of the middle row and compare it with its neighbors above and below. See Tipping Point & Peak.

Hint: Count Negative Numbers in a Sorted Matrix

Start at the bottom-left corner. A negative cell means the whole rest of its row is negative, so count them and move up; otherwise move right. That is one staircase walk.

Key takeaways#

Check how the matrix is sorted before choosing a technique.
Fully chained rows: flatten, and map i to (i // n, i % n).
Rows and columns sorted separately: staircase from the top-right, $O(m + n)$ .
Kth smallest: binary search the value and count cells ≤ guess.
Think of the 1D pattern first, then map it onto the grid.

Core idea#

A sorted matrix is usually a 1D pattern in disguise. First work out how it is sorted. Then flatten it, walk a staircase from a corner, or binary search the values.

Pick the technique#

Matrix shape	Technique	Cost
Each row starts after the previous row ends	flatten to one sorted array	$O(\log(mn))$
Rows and columns sorted separately	staircase from top-right	$O(m + n)$
Rows and columns sorted, Kth smallest	binary search the value	$O((m + n) \log(\text{range}))$

Key points#

Flatten: index i lives at matrix[i // n][i % n]. Search 0 .. m*n - 1.
Staircase: start at the top-right cell. Everything left is smaller. Everything below is larger.
Cell bigger than the target: move left. Cell smaller: move down.
Each staircase step removes a row or a column.
Binary searching every row instead costs $O(m \log n)$ .
Kth smallest: count cells ≤ a guess with the staircase walk. Find the first value with count ≥ k.
2D peak: binary search over rows with the maximum of the middle row.

Template#

row, col = 0, last column              # top-right corner
while row is in range and col >= 0:
    if matrix[row][col] > target: col -= 1   # rest of column bigger
    elif matrix[row][col] < target: row += 1 # rest of row smaller
    else: return True
return False

Pitfalls#

Calling the staircase logarithmic. It is $O(m + n)$ .
The wrong corner. Top-left and bottom-right give no decision. Use top-right or bottom-left.
Flattening a matrix whose rows do not chain.
Extra min/max guards on row indices. They hide bugs.

Say it in an interview#

"First, check how the matrix is sorted."
"From the top-right, each comparison removes a row or a column, so it is O(m + n)."