01 Matrix

35 min · updateMatrix()

A city map marks safe stations with 0 and every other block with 1. For each block, you want one number: how many street steps away is the nearest station?

Given a rectangular binary matrix mat, return a matrix of the same shape. The value at each position must be its shortest distance to any 0. One step moves up, down, left, or right; diagonal movement doesn't count.

Ground rules:

Every input contains at least one 0, so each distance exists.
A 0 has distance 0 to itself.
Return the completed distance matrix; don't return paths or coordinates.

Constraints

1 <= mat.length <= 200
1 <= mat[i].length <= 200
Every row has the same length.
Every cell is 0 or 1.
At least one cell is 0.

Hints

One search per cell

Starting a fresh breadth-first search from every 1 works. On a 200 × 200 matrix, how many times might those searches walk across the same cells?

Reverse the question

Instead of asking each 1 to find a 0, let every 0 spread its distance outward at the same time. The first wave to reach a cell must come from its nearest zero.

Seed before you search

Put every zero in one queue with distance 0. Mark every other distance as unknown, then assign a neighbor only the first time a wave reaches it.

Follow-up

Suppose some cells become walls that no path may cross. What would you change, and what value would you return for a cell that can't reach any zero?

Visible cases

Examples

Example 1

ready

Input: mat = [ [ 0, 0, 0 ], [ 0, 1, 0 ], [ 0, 0, 0 ] ]
Expected: [ [ 0, 0, 0 ], [ 0, 1, 0 ], [ 0, 0, 0 ] ]
Why: the center is one step from a zero; zero cells keep distance 0

Example 2

ready

Input: mat = [ [ 0, 1, 1, 1, 1 ], [ 1, 1, 1, 1, 1 ], [ 1, 1, 1, 1, 1 ], [ 1, 1, 1, 1, 1 ], [ 1, 1, 1, 1, 1 ] ]
Expected: [ [ 0, 1, 2, 3, 4 ], [ 1, 2, 3, 4, 5 ], [ 2, 3, 4, 5, 6 ], [ 3, 4, 5, 6, 7 ], [ 4, 5, 6, 7, 8 ] ]
Why: distance grows one layer at a time from the lone zero in the corner

Example 3

ready

Input: mat = [ [ 1, 0 ], [ 1, 1 ] ]
Expected: [ [ 1, 0 ], [ 2, 1 ] ]
Why: the lower-left cell is two orthogonal steps from the zero

Interview signal

Asked at

AmazonGoogleMeta

BFS from every one

Grid BFS shortest path is breadth-first search over cells, where each legal move is one edge. Because every edge costs one step, the first zero reached from a starting cell is a nearest zero.

Run that search separately for every 1 and write each answer into a result matrix.

from collections import deque
 
 
def update_matrix(mat: list[list[int]]) -> list[list[int]]:
    rows, cols = len(mat), len(mat[0])
    answer = [[0] * cols for _ in range(rows)]
 
    for start_row in range(rows):
        for start_col in range(cols):
            if mat[start_row][start_col] == 0:
                continue
 
            queue = deque([(start_row, start_col, 0)])
            seen = {(start_row, start_col)}
            while queue:
                row, col, distance = queue.popleft()
                if mat[row][col] == 0:
                    answer[start_row][start_col] = distance
                    break
                for dr, dc in ((1, 0), (-1, 0), (0, 1), (0, -1)):
                    nr, nc = row + dr, col + dc
                    if 0 <= nr < rows and 0 <= nc < cols and (nr, nc) not in seen:
                        seen.add((nr, nc))
                        queue.append((nr, nc, distance + 1))
 
    return answer

Time O((mn)²) — as many as mn searches may each cross mn cells. Space O(mn) for one search queue and its visited set.

Multi-source BFS — let every zero compete

The repeated searches overlap almost completely. Reverse their direction with multi-source BFS, one breadth-first search seeded from every zero at once.

Initialize zero cells to distance 0 and every other cell to -1, meaning unvisited. When a queue cell reaches an unvisited neighbor, its distance is the current distance plus one. BFS layers arrive in increasing distance order, so that first assignment is already optimal.

from collections import deque
 
 
def update_matrix(mat: list[list[int]]) -> list[list[int]]:
    rows, cols = len(mat), len(mat[0])
    distance = [[-1] * cols for _ in range(rows)]
    queue: deque[tuple[int, int]] = deque()
 
    for row in range(rows):
        for col in range(cols):
            if mat[row][col] == 0:
                distance[row][col] = 0
                queue.append((row, col))
 
    while queue:
        row, col = queue.popleft()
        for dr, dc in ((1, 0), (-1, 0), (0, 1), (0, -1)):
            nr, nc = row + dr, col + dc
            if 0 <= nr < rows and 0 <= nc < cols and distance[nr][nc] == -1:
                distance[nr][nc] = distance[row][col] + 1
                queue.append((nr, nc))
 
    return distance

The distance matrix doubles as the visited record. A cell enters the queue once, and the nearest zero claims it before any longer route can.

Time O(mn) — one visit per cell. Space O(mn) for the returned matrix and queue; the queue alone is also O(mn) in the widest layer.

Loading editor

Console

ready to run

Ready to run.

3 cases are queued.

Visible testcase

Case 1

Arguments · JSON arraymat: int[][]

Expected

Run: visible + custom · Submit: full suite