Searching a sorted matrix

Interviewers use two different kinds of sorted matrix, and they need different algorithms. In a fully sorted matrix, each row is sorted and each row starts after the previous row ends, so reading row by row gives one sorted list. You can treat it as a flat array of m x n elements and map index i to row i // n and column i % n.

In a row-and-column sorted matrix, every row is sorted left to right and every column top to bottom, but a row can start below where the previous one ended. Here you start at the top-right corner: if the value is too big, move left (the whole column below is bigger); if it is too small, move down (the whole row to the left is smaller). Each step removes a row or a column.

1. 1Fully sorted matrix: set lo = 0 and hi = m x n - 1 and run classic binary search.
2. 2For each mid, read matrix[mid // n][mid % n], where n is the number of columns. Compare and move lo or hi exactly as in a 1D array.
3. 3Row-and-column sorted matrix: start at row 0, column n - 1 (top-right). From there, going left decreases values and going down increases them, so every comparison has one clear direction.
4. 4If the value equals the target, done. If it is greater, the target cannot be in this column (everything below is even greater): move left. If it is smaller, the target cannot be in this row (everything to the left is even smaller): move down.
5. 5Stop when you leave the matrix: the target is not present. You take at most m + n steps.
6. 6For 'kth smallest' in a row-and-column sorted matrix, binary search on the value range and use the staircase walk to count entries <= mid.

def search_flat_matrix(matrix, target):
    """Rows sorted and each row starts after the previous row ends."""
    if not matrix or not matrix[0]:
        return False
    rows, cols = len(matrix), len(matrix[0])
    lo, hi = 0, rows * cols - 1        # treat the matrix as one sorted array
    while lo <= hi:
        mid = (lo + hi) // 2
        value = matrix[mid // cols][mid % cols]
        if value == target:
            return True
        if value < target:
            lo = mid + 1
        else:
            hi = mid - 1
    return False

def search_staircase(matrix, target):
    """Every row and every column sorted ascending. O(rows + cols)."""
    if not matrix or not matrix[0]:
        return False
    r, c = 0, len(matrix[0]) - 1       # start at the top-right corner
    while r < len(matrix) and c >= 0:
        if matrix[r][c] == target:
            return True
        if matrix[r][c] > target:
            c -= 1                     # whole column below is even bigger
        else:
            r += 1                     # whole row to the left is even smaller
    return False

search_flat_matrix is classic binary search over virtual indexes 0 to m x n - 1, using // and % to convert to row and column: O(log(m x n)) time. search_staircase removes one row or one column per step: O(m + n) time. Both use O(1) extra space.

Cost: Fully sorted: O(log(m x n)) = O(log m + log n). Row-and-column sorted: O(m + n). Both O(1) space.