542 01-矩阵

标签: BFS, 图, 广度优先

题目

题目链接：542 01-矩阵

给定一个由 0 和 1 组成的矩阵，找出每个元素到最近的 0 的距离。

两个相邻元素间的距离为 1 。

示例 1：

输入：
    [[0,0,0],
     [0,1,0],
     [0,0,0]]
输出：
    [[0,0,0],
     [0,1,0],
     [0,0,0]] 示例 2：

输入：
    [[0,0,0],
     [0,1,0],
     [1,1,1]]

输出：
    [[0,0,0],
     [0,1,0],
     [1,2,1]]

提示：

给定矩阵的元素个数不超过 10000。
给定矩阵中至少有一个元素是 0。
矩阵中的元素只在四个方向上相邻: 上、下、左、右。

思路

1. 依次遍历矩阵中为 1 的元素，然后对该元素进行广度优先搜索，一层一层向外搜索，找到离该元素最近的 0 元素，并返回二者之间的距离，即为「最近的 0 的距离」。

2. 因为要求「最近的 0 的距离」，所以思路应该是广度优先的层次遍历，使用队列来实现：

代码

class Solution:
    def updateMatrix(self, matrix: List[List[int]]) -> List[List[int]]:
        import queue
        def bfs(row, col):
            q = queue.Queue()
            q.put((row, col))

            # 记录访问过的元素
            visited = set()
            # 记录层次
            level = 0

            while q.qsize() > 0:
                sz = q.qsize()
                # 按层遍历
                for i in range(sz):
                    curX, curY = q.get()
                    # 找到最近的 0，level 即为距离
                    if matrix[curX][curY] == 0:
                        return level
                    for dx, dy in [[-1,0], [1,0], [0,-1], [0,1]]:
                        if 0 <= curX+dx < len(matrix) and 0 <= curY+dy < len(matrix[0]) and (curX+dx, curY+dy) not in visited:
                            q.put((curX+dx, curY+dy))
                            visited.add((curX+dx, curY+dy))
                level += 1
            return level

        # 初始化
        res = [[0] * len(matrix[0]) for i in range(len(matrix))]
        
        # 遍历矩阵
        for r in range(len(matrix)):
            for c in range(len(matrix[0])):
                # 找到 1 的元素，进行 BFS 搜索
                if matrix[r][c] == 1:
                    res[r][c] = bfs(r,c)
        return res