542. 01 矩阵
题目描述
输入一个 m x n 的矩阵 mat,每个格子的值为 0 或 1。
在矩阵中,每次只能走到上、下、左、右相邻的格子,每走一步距离加 1。对每个格子,求它走到离它最近的值为 0 的格子所需的最少步数;值为 0 的格子自身距离为 0。
返回一个同样大小的矩阵,第 i 行第 j 列存放格子 (i, j) 的这个距离。
输入:mat = [[0,0,0],[0,1,0],[0,0,0]] 输出:[[0,0,0],[0,1,0],[0,0,0]]
解释:唯一的 1 在中心,它四周都是 0,距离为 1。
输入: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.lengthn == mat[i].length1 <= m, n <= 10^41 <= m * n <= 10^4mat[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 distclass Solution {
public int[][] updateMatrix(int[][] mat) {
int m = mat.length, n = mat[0].length;
int[][] dist = new int[m][n];
ArrayDeque<int[]> q = new ArrayDeque<>();
for (int i = 0; i < m; i++) {
for (int j = 0; j < n; j++) {
dist[i][j] = -1;
if (mat[i][j] == 0) {
dist[i][j] = 0;
q.add(new int[]{i, j});
}
}
}
int[][] dirs = {{1, 0}, {-1, 0}, {0, 1}, {0, -1}};
while (!q.isEmpty()) {
int[] cur = q.poll();
for (int[] d : dirs) {
int x = cur[0] + d[0], y = cur[1] + d[1];
if (x >= 0 && x < m && y >= 0 && y < n && dist[x][y] == -1) {
dist[x][y] = dist[cur[0]][cur[1]] + 1;
q.add(new int[]{x, y});
}
}
}
return dist;
}
}class Solution {
public:
vector<vector<int>> updateMatrix(vector<vector<int>>& mat) {
int m = mat.size(), n = mat[0].size();
vector<vector<int>> dist(m, vector<int>(n, -1));
queue<pair<int, int>> q;
for (int i = 0; i < m; i++) {
for (int j = 0; j < n; j++) {
if (mat[i][j] == 0) {
dist[i][j] = 0;
q.push({i, j});
}
}
}
int dx[4] = {1, -1, 0, 0}, dy[4] = {0, 0, 1, -1};
while (!q.empty()) {
auto [i, j] = q.front();
q.pop();
for (int d = 0; d < 4; d++) {
int x = i + dx[d], y = j + dy[d];
if (x >= 0 && x < m && y >= 0 && y < n && dist[x][y] == -1) {
dist[x][y] = dist[i][j] + 1;
q.push({x, y});
}
}
}
return dist;
}
};func updateMatrix(mat [][]int) [][]int {
m, n := len(mat), len(mat[0])
dist := make([][]int, m)
q := [][2]int{}
for i := 0; i < m; i++ {
dist[i] = make([]int, n)
for j := 0; j < n; j++ {
dist[i][j] = -1
if mat[i][j] == 0 {
dist[i][j] = 0
q = append(q, [2]int{i, j})
}
}
}
dirs := [][2]int{{1, 0}, {-1, 0}, {0, 1}, {0, -1}}
for len(q) > 0 {
c := q[0]
q = q[1:]
for _, d := range dirs {
x, y := c[0]+d[0], c[1]+d[1]
if x >= 0 && x < m && y >= 0 && y < n && dist[x][y] == -1 {
dist[x][y] = dist[c[0]][c[1]] + 1
q = append(q, [2]int{x, y})
}
}
}
return dist
}var updateMatrix = function(mat) {
const m = mat.length, n = mat[0].length;
const dist = Array.from({ length: m }, () => new Array(n).fill(-1));
const q = [];
for (let i = 0; i < m; i++) {
for (let j = 0; j < n; j++) {
if (mat[i][j] === 0) {
dist[i][j] = 0;
q.push([i, j]);
}
}
}
for (let h = 0; h < q.length; h++) {
const [i, j] = q[h];
for (const [x, y] of [[i + 1, j], [i - 1, j], [i, j + 1], [i, j - 1]]) {
if (x >= 0 && x < m && y >= 0 && y < n && dist[x][y] === -1) {
dist[x][y] = dist[i][j] + 1;
q.push([x, y]);
}
}
}
return dist;
};function updateMatrix(mat: number[][]): number[][] {
const m = mat.length, n = mat[0].length;
const dist = Array.from({ length: m }, () => new Array(n).fill(-1));
const q: number[][] = [];
for (let i = 0; i < m; i++) {
for (let j = 0; j < n; j++) {
if (mat[i][j] === 0) {
dist[i][j] = 0;
q.push([i, j]);
}
}
}
for (let h = 0; h < q.length; h++) {
const [i, j] = q[h];
for (const [x, y] of [[i + 1, j], [i - 1, j], [i, j + 1], [i, j - 1]]) {
if (x >= 0 && x < m && y >= 0 && y < n && dist[x][y] === -1) {
dist[x][y] = dist[i][j] + 1;
q.push([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 distclass Solution {
public int[][] updateMatrix(int[][] mat) {
int m = mat.length, n = mat[0].length;
int INF = m + n;
int[][] dist = new int[m][n];
for (int i = 0; i < m; i++)
for (int j = 0; j < n; j++)
dist[i][j] = mat[i][j] == 0 ? 0 : INF;
for (int i = 0; i < m; i++) {
for (int j = 0; j < n; j++) {
if (i > 0) dist[i][j] = Math.min(dist[i][j], dist[i - 1][j] + 1);
if (j > 0) dist[i][j] = Math.min(dist[i][j], dist[i][j - 1] + 1);
}
}
for (int i = m - 1; i >= 0; i--) {
for (int j = n - 1; j >= 0; j--) {
if (i < m - 1) dist[i][j] = Math.min(dist[i][j], dist[i + 1][j] + 1);
if (j < n - 1) dist[i][j] = Math.min(dist[i][j], dist[i][j + 1] + 1);
}
}
return dist;
}
}class Solution {
public:
vector<vector<int>> updateMatrix(vector<vector<int>>& mat) {
int m = mat.size(), n = mat[0].size();
int INF = m + n;
vector<vector<int>> dist(m, vector<int>(n));
for (int i = 0; i < m; i++)
for (int j = 0; j < n; j++)
dist[i][j] = mat[i][j] == 0 ? 0 : INF;
for (int i = 0; i < m; i++) {
for (int j = 0; j < n; j++) {
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 (int i = m - 1; i >= 0; i--) {
for (int j = n - 1; j >= 0; j--) {
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;
}
};func updateMatrix(mat [][]int) [][]int {
m, n := len(mat), len(mat[0])
INF := m + n
dist := make([][]int, m)
for i := 0; i < m; i++ {
dist[i] = make([]int, n)
for j := 0; j < n; j++ {
if mat[i][j] != 0 {
dist[i][j] = INF
}
}
}
for i := 0; i < m; i++ {
for j := 0; j < n; j++ {
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 := m - 1; i >= 0; i-- {
for j := n - 1; j >= 0; j-- {
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
}var updateMatrix = function(mat) {
const m = mat.length, n = mat[0].length;
const INF = m + n;
const dist = mat.map(row => row.map(v => v === 0 ? 0 : INF));
for (let i = 0; i < m; i++) {
for (let j = 0; j < n; j++) {
if (i > 0) dist[i][j] = Math.min(dist[i][j], dist[i - 1][j] + 1);
if (j > 0) dist[i][j] = Math.min(dist[i][j], dist[i][j - 1] + 1);
}
}
for (let i = m - 1; i >= 0; i--) {
for (let j = n - 1; j >= 0; j--) {
if (i < m - 1) dist[i][j] = Math.min(dist[i][j], dist[i + 1][j] + 1);
if (j < n - 1) dist[i][j] = Math.min(dist[i][j], dist[i][j + 1] + 1);
}
}
return dist;
};function updateMatrix(mat: number[][]): number[][] {
const m = mat.length, n = mat[0].length;
const INF = m + n;
const dist = mat.map(row => row.map(v => v === 0 ? 0 : INF));
for (let i = 0; i < m; i++) {
for (let j = 0; j < n; j++) {
if (i > 0) dist[i][j] = Math.min(dist[i][j], dist[i - 1][j] + 1);
if (j > 0) dist[i][j] = Math.min(dist[i][j], dist[i][j - 1] + 1);
}
}
for (let i = m - 1; i >= 0; i--) {
for (let j = n - 1; j >= 0; j--) {
if (i < m - 1) dist[i][j] = Math.min(dist[i][j], dist[i + 1][j] + 1);
if (j < n - 1) dist[i][j] = Math.min(dist[i][j], dist[i][j + 1] + 1);
}
}
return dist;
}