LeetCode 542: 01 Matrix

Problem Restatement

We are given an m x n binary matrix mat.

Each cell contains either 0 or 1.

For every cell, we need to return the distance to the nearest cell containing 0.

The distance between two cells sharing an edge is 1. So movement is allowed only in four directions:

up, down, left, right

Diagonal movement is not allowed.

The problem guarantees that the matrix contains at least one 0. The official statement asks for the distance of the nearest 0 for each cell, with edge-adjacent cells having distance 1.

Input and Output

Function shape:

def updateMatrix(mat: list[list[int]]) -> list[list[int]]:
    ...

Examples

Consider:

mat = [
    [0, 0, 0],
    [0, 1, 0],
    [0, 0, 0],
]

The center cell is 1.

Its nearest zero is one step away in any direction.

So the output is:

[
    [0, 0, 0],
    [0, 1, 0],
    [0, 0, 0],
]

Another example:

mat = [
    [0, 0, 0],
    [0, 1, 0],
    [1, 1, 1],
]

The nearest-zero distances are:

[
    [0, 0, 0],
    [0, 1, 0],
    [1, 2, 1],
]

The bottom middle cell has distance 2, because it must move two edge steps to reach a zero.

First Thought: BFS From Each One

A direct idea is to process every 1 cell independently.

For each 1, run BFS until we find a 0.

This gives the correct distance for that cell.

But if there are many 1 cells, this repeats almost the same search many times.

In the worst case, this can become too slow because each cell may start its own BFS.

Key Insight

Instead of asking:

For each 1, where is the nearest 0?

we reverse the direction:

Starting from all 0 cells, how far does distance spread?

This is multi-source BFS.

All zero cells start in the queue with distance 0.

Then BFS expands outward one layer at a time.

The first time we reach a 1 cell, we have reached it by the shortest possible path from any zero.

That first distance is the answer for that cell.

Algorithm

Create a result matrix dist.

Initialize all cells to -1, meaning unvisited.

Then:

1. Scan the matrix.
2. For every cell with value 0:
   1. Set its distance to 0.
   2. Add it to the BFS queue.
3. Run BFS from all zero cells at once.
4. For each popped cell, check its four neighbors.
5. If a neighbor is inside the matrix and unvisited:
   1. Set its distance to current distance plus 1.
   2. Add it to the queue.
6. Return dist.

Correctness

All zero cells are placed into the BFS queue with distance 0. This is correct because the distance from a zero cell to the nearest zero is 0.

BFS processes cells in increasing distance order. Since all zeros are inserted before the search begins, the BFS expands simultaneously from every zero cell.

When the algorithm first visits a cell (r, c), it must have reached it through the shortest path from one of the zero cells. If a shorter path existed, BFS would have processed that shorter path earlier and visited (r, c) sooner.

The algorithm assigns a distance only once, at the first visit. Therefore, each assigned distance is the minimum distance to any zero.

Since BFS explores through all four edge-adjacent directions, it considers exactly the allowed moves.

Thus the returned matrix contains the nearest-zero distance for every cell.

Complexity

Let:

m = number of rows
n = number of columns

Each cell checks four neighbors, which is constant work.

Implementation

from collections import deque

class Solution:
    def updateMatrix(self, mat: list[list[int]]) -> list[list[int]]:
        rows = len(mat)
        cols = len(mat[0])

        dist = [[-1] * cols for _ in range(rows)]
        queue = deque()

        for r in range(rows):
            for c in range(cols):
                if mat[r][c] == 0:
                    dist[r][c] = 0
                    queue.append((r, c))

        directions = [
            (-1, 0),
            (1, 0),
            (0, -1),
            (0, 1),
        ]

        while queue:
            r, c = queue.popleft()

            for dr, dc in directions:
                nr = r + dr
                nc = c + dc

                if 0 <= nr < rows and 0 <= nc < cols and dist[nr][nc] == -1:
                    dist[nr][nc] = dist[r][c] + 1
                    queue.append((nr, nc))

        return dist

Code Explanation

We create the distance matrix:

dist = [[-1] * cols for _ in range(rows)]

A value of -1 means the cell has not been assigned a distance yet.

Then we put every zero cell into the queue:

if mat[r][c] == 0:
    dist[r][c] = 0
    queue.append((r, c))

This is the multi-source part of the BFS.

All zeros are sources.

The directions are the four allowed moves:

directions = [
    (-1, 0),
    (1, 0),
    (0, -1),
    (0, 1),
]

During BFS, for each neighbor, we check whether it is inside the matrix and unvisited:

if 0 <= nr < rows and 0 <= nc < cols and dist[nr][nc] == -1:

If so, its nearest zero distance is one more than the current cell’s distance:

dist[nr][nc] = dist[r][c] + 1

Then we add it to the queue so its neighbors can be processed later.

Dynamic Programming Version

There is also a two-pass dynamic programming solution.

The idea is to relax distances from top-left to bottom-right, then from bottom-right to top-left.

class Solution:
    def updateMatrix(self, mat: list[list[int]]) -> list[list[int]]:
        rows = len(mat)
        cols = len(mat[0])
        inf = rows + cols + 1

        dist = [[inf] * cols for _ in range(rows)]

        for r in range(rows):
            for c in range(cols):
                if mat[r][c] == 0:
                    dist[r][c] = 0
                else:
                    if r > 0:
                        dist[r][c] = min(dist[r][c], dist[r - 1][c] + 1)
                    if c > 0:
                        dist[r][c] = min(dist[r][c], dist[r][c - 1] + 1)

        for r in range(rows - 1, -1, -1):
            for c in range(cols - 1, -1, -1):
                if r + 1 < rows:
                    dist[r][c] = min(dist[r][c], dist[r + 1][c] + 1)
                if c + 1 < cols:
                    dist[r][c] = min(dist[r][c], dist[r][c + 1] + 1)

        return dist

The BFS version is usually easier to reason about because it directly computes shortest distance from all zero cells.

Testing

def run_tests():
    s = Solution()

    assert s.updateMatrix([
        [0, 0, 0],
        [0, 1, 0],
        [0, 0, 0],
    ]) == [
        [0, 0, 0],
        [0, 1, 0],
        [0, 0, 0],
    ]

    assert s.updateMatrix([
        [0, 0, 0],
        [0, 1, 0],
        [1, 1, 1],
    ]) == [
        [0, 0, 0],
        [0, 1, 0],
        [1, 2, 1],
    ]

    assert s.updateMatrix([
        [0],
    ]) == [
        [0],
    ]

    assert s.updateMatrix([
        [1, 1, 1],
        [1, 0, 1],
        [1, 1, 1],
    ]) == [
        [2, 1, 2],
        [1, 0, 1],
        [2, 1, 2],
    ]

    print("all tests passed")

run_tests()