给定一个由 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到0的最短距离,用BFS。
from collections import deque
class Solution(object):
def updateMatrix(self, matrix):
"""
:type matrix: List[List[int]]
:rtype: List[List[int]]
"""
if not matrix or not matrix[0]:
return matrix
m, n = len(matrix), len(matrix[0])
res = matrix[:]
q = deque()
for i in range(m):
for j in range(n):
if matrix[i][j] == 0:
q.append([[i, j], 0])
dx = [1, -1, 0, 0]
dy = [0, 0, 1, -1]
visited = [[0 for _ in range(n + 1)] for _ in range(m + 1)]
while q:
tmp, distance = q.popleft()
x0, y0 = tmp[0], tmp[1]
if matrix[x0][y0] == 1:
res[x0][y0] = distance
for k in range(4):
x = x0 + dx[k]
y = y0 + dy[k]
if 0 <= x < m and 0 <= y < n and visited[x][y] != 1:
q.append([[x, y], distance + 1])
visited[x][y] = 1
return res