用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

运行结果

542. 01 Matrix刷题笔记_双端队列


但是运行结果还是比较慢,究其原因,第一个是两个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))

542. 01 Matrix刷题笔记_双端队列_02