题目
给定一个由 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。
- 矩阵中的元素只在四个方向上相邻: 上、下、左、右。
解题思路
一般求最短路的问题都可以用广度优先搜索(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