[LeetCode]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

注意:

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

解题思路

一般求最短路的问题都可以用广度优先搜索(BFS)来解决。
首先将矩阵中所有的0想象成一个超级源点，第一步先将所有的0加入队列，所以最近的0就是它们本身，也即距离为0，然后将它们加入已访问集合中，当队列不为空时，每次从队列中弹出一个元素，判断其他四个方向的某元素是否有未访问过的，若未访问过，则该方向上的某元素到最近的0的距离就是当前元素到最近的0的距离加1，并将它加入已访问集合和队列中。

代码

class Solution:
    def updateMatrix(self, matrix: List[List[int]]) -> List[List[int]]:
        m = len(matrix)
        n = len(matrix[0])
        directions = [(-1,0), (1,0), (0,-1), (0,1)]
        queue = []
        visited = set() # 已访问元素集合
        for i in range(m):
            for j in range(n):
                if matrix[i][j] == 0:
                    queue.append((i,j))
                    visited.add((i,j))
        while queue:
            i, j = queue.pop(0)
            for dx, dy in directions:
                if 0<=i+dx<m and 0<=j+dy<n and (i+dx,j+dy) not in visited:
                    matrix[i+dx][j+dy] = matrix[i][j] + 1
                    visited.add((i+dx,j+dy))
                    queue.append((i+dx,j+dy))
        return matrix