542. 01 Matrix (Medium) - TengnanYao/daily_leetcode GitHub Wiki
class Solution:
def updateMatrix(self, mat: List[List[int]]) -> List[List[int]]:
# bfs
m, n = len(mat), len(mat[0])
for i in range(m):
for j in range(n):
if mat[i][j] != 0:
q = [(i, j)]
count = 0
found = False
seen = set()
while q and not found:
temp = []
count += 1
for a, b in q:
for dx, dy in [(1, 0), (0, 1), (-1, 0), (0, -1)]:
x, y = a + dx, b + dy
if 0 <= x < m and 0 <= y < n and (x, y) not in seen:
seen.add((x, y))
if mat[x][y] == 0:
mat[i][j] = count
found = True
break
temp.append((x, y))
q = temp
return mat
# dp
m, n = len(mat), len(mat[0])
dp = [[inf] * n for _ in range(m)]
for i in range(m):
for j in range(n):
if mat[i][j] == 1:
if i > 0:
dp[i][j] = min(dp[i - 1][j] + 1, dp[i][j])
if j > 0:
dp[i][j] = min(dp[i][j - 1] + 1, dp[i][j])
else:
dp[i][j] = 0
for i in range(m - 1, -1, -1):
for j in range(n - 1, -1, -1):
if mat[i][j] == 1:
if i < m - 1:
dp[i][j] = min(dp[i + 1][j] + 1, dp[i][j])
if j < n - 1:
dp[i][j] = min(dp[i][j + 1] + 1, dp[i][j])
return dp