Problem Statement

Given an m x n binary matrix mat, return the distance of the nearest 0 for each cell.

The distance between two adjacent cells is 1.

Input: mat = [[0,0,0],[0,1,0],[1,1,1]]
Output: [[0,0,0],[0,1,0],[1,2,1]]

BFS: Iterative

from collections import deque
class Solution:
    def updateMatrix(self, mat: List[List[int]]) -> List[List[int]]:
        ROWS, COLS = len(mat), len(mat[0])
        """
            * we need to find the distance to the nearest 0
            * we know that m >= n >= 1 ~ no need to check dims
            * add zeros to queue
            * we can change mat
        """
        queue = deque()
        # Add locations of zero cells to queue with distance of 0
        # T: O(N*M)
        # S: O(N*M)
        for r in range(ROWS):
            for c in range(COLS):
                if mat[r][c] == 0:
                    # (row, col, distance)
                    queue.append((r, c, 0))
                else:
                    mat[r][c] = 10_001
        """
            [
                [0,0,0],
                [0,1,0],
                [1,2,1]
            ]
    queue = []

    dist = 2
        """
        # T: O(N*M)
        # S: O(N*M)
        while queue:
            for _ in range(len(queue)):
                ro, co, dist = queue.popleft()
                dist += 1

                for dr, dc in [(-1,0),(0,1),(1,0),(0,-1)]:
                    r = ro + dr
                    c = co + dc

                    if 0 <= r < ROWS and 0 <= c < COLS and \
                        mat[r][c] > mat[ro][co] + 1:
                        mat[r][c] = dist
                        queue.append((r, c, dist))
        return mat

Leetcode 542 01 Matrix