01 Matrix

The problem description can be found in LeetCode.

Approach

This is another example of an implicit Graph, as the Flood Fill problem we saw earlier.

The problem looks like a usual BFS one, since it would be enough to search at the nearest 0 if the cell is not equal to 0.

However, if we imagine an input like:

1 1 1 1
1 1 1 1
1 1 1 1
1 1 1 0

We would need to make a BFS $(rows\ast cols)-1$ times, which is $O(row{s}^{2}col{s}^2)$.

A smarter way would be to build the result matrix from scratch. We can initialize our queue with only coordinates of 0 values. Then, we can perform a BFS algorithm to update the distances of the cells that are not equal to 0.

Solution

from collections import deque
from typing import Deque, Iterator, List, Set, Tuple


class Solution:
    def updateMatrix(self, mat: List[List[int]]) -> List[List[int]]:
        # not ideal, but just to avoid to have too many function parameters
        self.mat = mat
        self.rows, self.cols = len(mat), len(mat[0])

        visited = set()
        queue: Deque[Tuple[int, int]] = deque()

        for row in range(self.rows):
            for col in range(self.cols):
                if self.mat[row][col] == 0:
                    visited.add((row, col))
                    queue.append((row, col))

        self.bfs(row, col, queue, visited)

        return self.mat

    def bfs(
        self,
        row: int,
        col: int,
        queue: Deque[Tuple[int, int]],
        visited: Set[Tuple[int, int]],
    ) -> None:
        distance = 0
        while queue:
            for _ in range(len(queue)):
                row, col = queue.popleft()
                self.mat[row][col] = distance

                for nbr_row, nbr_col in self.get_neighbors(row, col):
                    if (nbr_row, nbr_col) in visited:
                        continue

                    visited.add((nbr_row, nbr_col))
                    queue.append((nbr_row, nbr_col))

            distance += 1

    def get_neighbors(self, row: int, col: int) -> Iterator[Tuple[int, int]]:
        deltas = [
            (-1, 0),  # up
            (+1, 0),  # down
            (0, -1),  # left
            (0, +1),  # right
        ]
        for dt_row, dt_col in deltas:
            nbr_row = row + dt_row
            nbr_col = col + dt_col

            if (
                0 <= nbr_row < self.rows
                and 0 <= nbr_col < self.cols
                and self.mat[nbr_row][nbr_col] != 0
            ):
                yield nbr_row, nbr_col

Complexity

• Time: $O(mn)$
• Space: $O(mn)$

Where $m$ and $n$ are the number of rows and columns of the input matrix, respectively.