[LeetCods] 01Matrix (Python)

 

 

from collections import deque

class Solution:
    def updateMatrix(self, mat: List[List[int]]) -> List[List[int]]:
        row, col = len(mat), len(mat[0])
        dx = [0, 0, 1, -1]
        dy = [1, -1, 0, 0]
        INF = 1e9
        q = deque()
        
        for i in range(row):
            for j in range(col):
                if mat[i][j] == 0:
                    q.append((i, j))
                else : 
                    mat[i][j] = INF
                    
        while q:
            x, y = q.popleft()
            for i in range(4):
                nx = x + dx[i]
                ny = y + dy[i]
                z = mat[x][y] + 1
                if 0 <= nx < row and 0 <= ny < col and mat[nx][ny] > z:
                    mat[nx][ny] = z
                    q.append((nx, ny))
        return mat