LeetCode-Python-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

注意:

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

思路：

找所有1到0的最短距离，用BFS。

from collections import deque
class Solution(object):
    def updateMatrix(self, matrix):
        """
        :type matrix: List[List[int]]
        :rtype: List[List[int]]
        """
        if not matrix or not matrix[0]:
            return matrix
        m, n = len(matrix), len(matrix[0])
        res = matrix[:]
        
        q = deque()
               
        for i in range(m):
            for j in range(n):
                if matrix[i][j] == 0:
                    q.append([[i, j], 0])

        dx = [1, -1, 0, 0]        
        dy = [0, 0, 1, -1]
        visited = [[0 for _ in range(n + 1)] for _ in range(m + 1)]

        while q:
            tmp, distance = q.popleft()
            x0, y0 = tmp[0], tmp[1]

            if matrix[x0][y0] == 1:
                res[x0][y0] = distance

            for k in range(4):
                x = x0 + dx[k]
                y = y0 + dy[k]

                if 0 <= x < m and 0 <= y < n and visited[x][y] != 1:
                    q.append([[x, y], distance + 1])
                    visited[x][y] = 1
        return res