用bfs
pycharm代码一稿
class Solution:
def updateMatrix(self, mat):
m = len(mat)
n = len(mat[0])
directions = [(-1,0),(1,0),(0,-1),(0,1)]
to_check2 = []
for i in range(m):
for j in range(n):
if mat[i][j] == 0:
continue
to_check1=[[i,j]]
checked=[]
sucess = 0
while not sucess:
while to_check1:
[x,y] = to_check1.pop(0)
if 0<=x<m and 0<=y<n and [x,y] not in checked:
if mat[x][y]==0:
sucess = 1
mat[i][j] = abs(i-x)+abs(j-y)
break
else:
to_check2+=[[x-1,y],[x+1,y],[x,y-1],[x,y+1]]
checked += [[x, y]]
to_check1=to_check2
to_check2=[]
return mat
mat = [[0,0,0],[0,1,0],[0,0,0]]
print(Solution().updateMatrix(mat))
mat = [[0,0,0],[0,1,0],[1,1,1]]
print(Solution().updateMatrix(mat))
运行超时,原因是从1开始搜索,改成从0搜索就好了,参考了这个大神的博客
m = len(mat)
n = len(mat[0])
directions = [(-1,0),(1,0),(0,-1),(0,1)]
q = []
for i in range(m):
for j in range(n):
if mat[i][j] == 0:
q += [[i,j]]
else:
mat[i][j] = 20000
while q:
[ix,iy] = q.pop(0)
for dx,dy in directions:
nx = ix+dx
ny = iy+dy
if 0<=nx<m and 0<=ny<n:
if mat[nx][ny] > mat[ix][iy]+1:
mat[nx][ny] = mat[ix][iy]+1
q += [[nx,ny]]
return mat
运行结果
但是运行结果还是比较慢,究其原因,第一个是两个if可以合并,第二个是列表可以改成collections.deque(),因为deque是双端队列,双端队列的append()和pop()的时间复杂度为O(1),而list的insert(0,value),append以及pop()的时间复杂度为O(n)。
最终代码
import collections
class Solution:
def updateMatrix(self, mat):
m = len(mat)
n = len(mat[0])
directions = [(-1,0),(1,0),(0,-1),(0,1)]
q = collections.deque()
for i in range(m):
for j in range(n):
if mat[i][j] == 0:
q.append((i,j))
else:
mat[i][j] = 20000
while q:
ix,iy = q.popleft()
for dx,dy in directions:
nx = ix+dx
ny = iy+dy
if 0<=nx<m and 0<=ny<n and mat[nx][ny] > mat[ix][iy]+1:
mat[nx][ny] = mat[ix][iy]+1
q.append((nx,ny))
return mat
mat = [[0,0,0],[0,1,0],[0,0,0]]
print(Solution().updateMatrix(mat))
mat = [[0,0,0],[0,1,0],[1,1,1]]
print(Solution().updateMatrix(mat))