leetcode542 01矩阵

import collections
class Solution:
    def updateMatrix(self, matrix: List[List[int]]) -> List[List[int]]:
        row = len(matrix)
        col = len(matrix[0])
        res = [[0] * col for _ in range(row)]
        visited = [[0] * col for _ in range(row)]
        deque = collections.deque()
        for i in range(row): 
            for j in range(col):
                if matrix[i][j] == 0:
                    deque.append((i, j))

        dirs = [(1, 0), (-1, 0), (0, 1), (0, -1)]
        step = 0
        while deque:
            size = len(deque)
            for i in range(size):
                i, j = deque.popleft()
                if matrix[i][j] == 1:
                    res[i][j] = step
                for dir in dirs:
                    newx, newy = i + dir[0], j + dir[1]
                    if newx < 0 or newx >= row or newy < 0 or newy >= col or visited[newx][newy] == 1:
                        continue
                    deque.append((newx, newy))
                    visited[newx][newy] = 1
            step += 1
        return res