LeetCode-542、01 矩阵-中等

LeetCode-542、01 矩阵-中等

给定一个由 0 和 1 组成的矩阵，找出每个元素到最近的 0 的距离。两个相邻元素间的距离为 1 。

示例 1：

输入:
0 0 0
0 1 0
0 0 0

输出:
0 0 0
0 1 0
0 0 0

示例 2：

输入:
0 0 0
0 1 0
1 1 1

输出:
0 0 0
0 1 0
1 2 1

注意：

给定矩阵的元素个数不超过 10000。
给定矩阵中至少有一个元素是 0。
矩阵中的元素只在四个方向上相邻：上、下、左、右。

 

代码：

class Solution:
    def updateMatrix(self, matrix: List[List[int]]) -> List[List[int]]:
        m, n = len(matrix), len(matrix[0])
        queue = []
        visit = set()
        for i in range(m):
            for j in range(n):
                if matrix[i][j] == 0:
                    queue.append((i, j))
                    visit.add((i, j))
        while queue:
            i, j = queue.pop(0)
            for x, y in [(i+1, j), (i-1, j), (i, j+1), (i, j-1)]:
                if 0 <= x < m and 0 <= y < n and (x, y) not in visit:
                    matrix[x][y] = matrix[i][j] + 1
                    visit.add((x, y))
                    queue.append((x, y))
        return matrix