跳至主要內容

0542. 01 矩阵

ITCharge大约 3 分钟

0542. 01 矩阵open in new window

  • 标签:广度优先搜索、数组、动态规划、矩阵
  • 难度:中等

题目链接

题目大意

描述:给定一个 mnm * n 大小的、由 01 组成的矩阵 matmat

要求:输出一个大小相同的矩阵 resres,其中 res[i][j]res[i][j] 表示对应位置元素(即 mat[i][j]mat[i][j])到最近的 00 的距离。

说明

  • 两个相邻元素间的距离为 11
  • m==mat.lengthm == mat.length
  • n==mat[i].lengthn == mat[i].length
  • 1m,n1041 \le m, n \le 10^4
  • 1mn1041 \le m * n \le 10^4
  • mat[i][j]===0mat[i][j] === 0 或者 mat[i][j]==1mat[i][j] == 1
  • matmat 中至少有一个 00

示例

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

解题思路

思路 1:广度优先搜索

题目要求的是每个 10的最短曼哈顿距离。

比较暴力的做法是,从每个 1 开始进行广度优先搜索,每一步累积距离,当搜索到第一个 0,就是离这个 1 最近的 0,我们更新对应 1 位置上的答案距离。然后从下一个 1 开始进行广度优先搜索。

这样做每次进行广度优先搜索的时间复杂度为 O(m×n)O(m \times n)。对于 m×nm \times n 个节点来说,每个节点可能都要进行一次广度优先搜索,总的时间复杂度为 O(m2×n2)O(m^2 \times n^2)。时间复杂度太高了。

我们可以换个角度:求每个 01 的最短曼哈顿距离(和求每个 10 是等价的)。

我们将所有值为 0 的元素位置保存到队列中,然后对所有值为 0 的元素开始进行广度优先搜索,每搜一步距离加 1,当每次搜索到 1 时,就可以得到 0 到这个 1 的最短距离,也就是当前离这个 1 最近的 0 的距离。

这样对于所有节点来说,总共需要进行一次广度优先搜索就可以了,时间复杂度为 O(m×n)O(m \times n)

具体步骤如下:

  1. 使用一个集合变量 visited 存储所有值为 0 的元素坐标。使用队列变量 queue 存储所有值为 0 的元素坐标。使用二维数组 res 存储对应位置元素(即 mat[i][j]mat[i][j])到最近的 00 的距离。
  2. 我们从所有为如果队列 queue 不为空,则从队列中依次取出值为 0 的元素坐标,遍历其上、下、左、右位置。
  3. 如果相邻区域未被访问过(说明遇到了值为 1 的元素),则更新相邻位置的距离值,并把相邻位置坐标加入队列 queue 和访问集合 visited 中。
  4. 继续执行 2 ~ 3 步,直到队列为空时,返回 res

思路 1:代码

import collections

class Solution:
    def updateMatrix(self, mat: List[List[int]]) -> List[List[int]]:
        rows, cols = len(mat), len(mat[0])
        res = [[0 for _ in range(cols)] for _ in range(rows)]
        visited = set()

        for i in range(rows):
            for j in range(cols):
                if mat[i][j] == 0:
                    visited.add((i, j))
        
        directions = {(1, 0), (-1, 0), (0, 1), (0, -1)}
        queue = collections.deque(visited)

        while queue:
            i, j = queue.popleft()
            for direction in directions:
                new_i = i + direction[0]
                new_j = j + direction[1]
                if 0 <= new_i < rows and 0 <= new_j < cols and (new_i, new_j) not in visited:
                    res[new_i][new_j] = res[i][j] + 1
                    queue.append((new_i, new_j))
                    visited.add((new_i, new_j))
        return res

思路 1:复杂度分析

  • 时间复杂度O(m×n)O(m \times n)
  • 空间复杂度O(m×n)O(m \times n)