在问 · 算法题解

542. 01 矩阵

中等图论多源广度优先搜索力扣原题 ↗

题目描述

输入一个 m x n 的矩阵 mat,每个格子的值为 0 或 1。

在矩阵中,每次只能走到上、下、左、右相邻的格子,每走一步距离加 1。对每个格子,求它走到离它最近的值为 0 的格子所需的最少步数;值为 0 的格子自身距离为 0。

返回一个同样大小的矩阵,第 i 行第 j 列存放格子 (i, j) 的这个距离。

示例 1
输入:mat = [[0,0,0],[0,1,0],[0,0,0]]
输出:[[0,0,0],[0,1,0],[0,0,0]]

解释:唯一的 1 在中心,它四周都是 0,距离为 1。

示例 2
输入:mat = [[0,0,0],[0,1,0],[1,1,1]]
输出:[[0,0,0],[0,1,0],[1,2,1]]

解释:底行两端的 1 与上方的 0 相邻,距离为 1;底行中间的 1 要走两步才到 0。

数据范围

  • m == mat.length
  • n == mat[i].length
  • 1 <= m, n <= 10^4
  • 1 <= m * n <= 10^4
  • mat[i][j] 只能是 0 或 1。
  • mat 中至少有一个 0。

解题思路:多源广度优先搜索

多源 BFS。每个 1 各自搜索最近的 0 会重复很多工作。反过来:把所有 0 同时作为起点放进队列,距离设为 0,其余格子标为未访问(-1)。弹出一个格子,把它未访问的邻居设为距离 +1 并入队。BFS 按层扩展,第一次到达某格子的就是离它最近的那个 0。

为什么正确

队列中格子的距离按入队顺序非递减,且相邻两个最多差 1;每个格子在第一次被设置时就是最终距离。

复杂度分析

时间 O(mn),每个格子入队一次;空间 O(mn)。

易错点

入队时就要设置距离(标记已访问),否则同一格子会被重复入队。从每个 1 单独 BFS 的做法最坏 O((mn)^2),会超时。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言,逻辑一致。

class Solution:
    def updateMatrix(self, mat: List[List[int]]) -> List[List[int]]:
        m, n = len(mat), len(mat[0])
        dist = [[-1] * n for _ in range(m)]
        q = collections.deque()
        for i in range(m):
            for j in range(n):
                if mat[i][j] == 0:
                    dist[i][j] = 0
                    q.append((i, j))
        while q:
            i, j = q.popleft()
            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 dist[x][y] == -1:
                    dist[x][y] = dist[i][j] + 1
                    q.append((x, y))
        return dist

载入编辑器,自己跑一遍 →

其他解法

多源广度优先搜索通用,可以推广到任意无权图;两遍动态规划利用网格结构,省掉队列。两者都是 O(m·n)。

解法 2:两遍动态规划 时间 O(m·n),空间 O(1)

0 的距离是 0,1 的距离先设为一个足够大的数。第一遍从左上到右下,用上方和左方格子的距离加 1 更新当前格。第二遍从右下到左上,用下方和右方格子的距离加 1 再更新。任一条最短路径都可以拆成先向左上、再向右下的两段这类单调路径,所以两遍就够。

不需要队列,除结果矩阵外只用常数空间,常数也更小。正确性依赖曼哈顿距离的单调路径性质,只适用于网格且边权相同的情形。

class Solution:
    def updateMatrix(self, mat: List[List[int]]) -> List[List[int]]:
        m, n = len(mat), len(mat[0])
        INF = m + n
        dist = [[0 if mat[i][j] == 0 else INF for j in range(n)] for i in range(m)]
        for i in range(m):
            for j in range(n):
                if i > 0:
                    dist[i][j] = min(dist[i][j], dist[i - 1][j] + 1)
                if j > 0:
                    dist[i][j] = min(dist[i][j], dist[i][j - 1] + 1)
        for i in range(m - 1, -1, -1):
            for j in range(n - 1, -1, -1):
                if i < m - 1:
                    dist[i][j] = min(dist[i][j], dist[i + 1][j] + 1)
                if j < n - 1:
                    dist[i][j] = min(dist[i][j], dist[i][j + 1] + 1)
        return dist

图论 专题的其他题目

查看 图论 全部题目 →