Search a 2D Matrix (LeetCode 74)

The problem

You get an R × C integer matrix with two properties: each row is sorted left to right, and the first number of each row is greater than the last number of the previous row. Given a target, return True if it is in the matrix and False otherwise, in O(log(R·C)) time.

Input:
 1   3   5   7
10  11  16  20
23  30  34  60
target = 3

Output: True

3 is in the first row.

How to think about it

Start with brute force. Check every cell: O(R·C). A better step is to binary-search each row, which is O(R log C). Better again: binary-search the first column to find the right row, then binary-search that row, for O(log R + log C). That is already optimal, but it’s two searches and two chances for off-by-one errors.

Read the matrix row by row. Because each row starts above where the previous one ended, reading row 0, then row 1, then row 2 gives a single sorted list: 1, 3, 5, 7, 10, 11, 16, 20, 23, …. The matrix is just a sorted array of length R·C written in rows.

The key insight: you don’t need to build that array, only to map indices. Position k in the flat list is at row k // C and column k % C. So run an ordinary binary search over k in [0, R·C − 1] and read matrix[mid // C][mid % C] each time. One search, the familiar template, O(1) extra space.

Which pattern is this? Plain binary search over a virtual index space. The tell is a 2-D structure that is fully sorted in reading order. Note that LeetCode 240 (rows and columns sorted separately, without the row-to-row guarantee) is a different problem, solved by walking from a corner.

Approach

1. Let R, C be the matrix dimensions. Set lo = 0, hi = R*C - 1.
2. While lo <= hi:
   • mid = (lo + hi) // 2, and v = matrix[mid // C][mid % C].
   • If v == target, return True.
   • If v < target, set lo = mid + 1; otherwise hi = mid - 1.
3. Return False.

Solution

def search_matrix(matrix, target):
    R, C = len(matrix), len(matrix[0]); lo, hi = 0, R*C - 1   # flat index range
    while lo <= hi:
        mid = (lo + hi) // 2; v = matrix[mid // C][mid % C]  # divide by columns, not rows
        if v == target: return True
        if v < target: lo = mid + 1
        else: hi = mid - 1
    return False
# search_matrix([[1,3,5,7],[10,11,16,20],[23,30,34,60]], 3) -> True

Dry run

On the example, with C = 4 and target = 3:

The answer is True. Searching for 13 instead visits 11, 23 and 16, then the range empties and it returns False.

Complexity

• Time: O(log(R·C)). One binary search over R·C positions, which equals O(log R + log C).
• Space: O(1). The flat array is never built.

Edge cases and pitfalls

• Divide by C, the number of columns, for both the row and the column. Using R is the classic bug and only “works” on square matrices.
• Single row or single column matrices are fine: the mapping still holds.
• Empty matrix ([] or [[]]) would crash on matrix[0] or give hi = -1. LeetCode guarantees at least one cell, but guard it if asked.
• Don’t use this on LeetCode 240. Without the row-to-row ordering, the flattened list isn’t sorted.

Related problems

• Binary Search (LeetCode 704): the exact template used here.
• Spiral Matrix (LeetCode 54): more practice converting between 2-D positions and a traversal order.
• Kth Smallest Element in a BST (LeetCode 230): another structure whose in-order reading is a sorted list.