01 Matrix

Problem Id: 542 Difficulty: Medium


Intuition

Solution


class Solution:
    def updateMatrix(self, matrix: List[List[int]]) -> List[List[int]]:
        # Edge cases
        if not matrix:
            return []

        # Prepare
        distance = 0
        dfs = []
        for i in range(len(matrix)):
            for j in range(len(matrix[0])):
                if matrix[i][j] == 0:
                    dfs.append((i, j))
        ans = [[-1] * len(matrix[0]) for _ in range(len(matrix))]

        # DFS mark distance
        while dfs:
            tmp = []
            for i, j in dfs:
                ans[i][j] = distance
            for i, j in dfs:
                for ni, nj in [(i - 1, j), (i + 1, j), (i, j - 1), (i, j + 1)]:
                    if ni < 0 or nj < 0 or ni >= len(matrix) or nj >= len(matrix[0]):
                        continue
                    if ans[ni][nj] != -1:
                        continue
                    tmp.append((ni, nj))
            dfs = tmp
            distance += 1
            tmp = []

        return ans