01 Matrix

Medium Problem 103 of 151 Java and Python

Problem

Given an m × n binary matrix mat, return a matrix of the same size where each cell holds the distance to the nearest 0. Moving to an adjacent cell (up, down, left, right) has distance 1.

Example 1

Input: mat = [[0, 0, 0], [0, 1, 0], [0, 0, 0]]
Output: [[0, 0, 0], [0, 1, 0], [0, 0, 0]]

Example 2

Input: mat = [[0, 0, 0], [0, 1, 0], [1, 1, 1]]
Output: [[0, 0, 0], [0, 1, 0], [1, 2, 1]]

Constraints

1 ≤ m · n ≤ 10⁴
There is at least one 0.

Explanation

Running a BFS from every 1 repeats a lot of work. Turn it around: start one BFS from all zeros at once. The first time the wave reaches a cell is its distance to the nearest zero.

Approach 1 · Naive

BFS from every 1

Time O((m · n)²)
Space O(m · n)

For each cell containing 1, run a BFS until it reaches a 0.

For each 1, BFS outwards with a fresh visited grid.
The level at which you first see a 0 is the distance.

Java

import java.util.*;

class Solution {
    public int[][] updateMatrix(int[][] mat) {
        int m = mat.length;
        int n = mat[0].length;
        int[][] result = new int[m][n];
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                if (mat[i][j] == 1) {
                    result[i][j] = nearestZero(mat, i, j);
                }
            }
        }
        return result;
    }

    private int nearestZero(int[][] mat, int sr, int sc) {
        int m = mat.length;
        int n = mat[0].length;
        boolean[][] seen = new boolean[m][n];
        Deque<int[]> queue = new ArrayDeque<>();
        queue.add(new int[] { sr, sc, 0 });
        seen[sr][sc] = true;
        int[][] dirs = { { 1, 0 }, { -1, 0 }, { 0, 1 }, { 0, -1 } };
        while (!queue.isEmpty()) {
            int[] cell = queue.poll();
            if (mat[cell[0]][cell[1]] == 0) {
                return cell[2];
            }
            for (int[] d : dirs) {
                int r = cell[0] + d[0];
                int c = cell[1] + d[1];
                if (r >= 0 && c >= 0 && r < m && c < n && !seen[r][c]) {
                    seen[r][c] = true;
                    queue.add(new int[] { r, c, cell[2] + 1 });
                }
            }
        }
        return -1;
    }
}

Python

from collections import deque

class Solution:
    def updateMatrix(self, mat: List[List[int]]) -> List[List[int]]:
        m, n = len(mat), len(mat[0])

        def nearest_zero(sr, sc):
            seen = {(sr, sc)}
            queue = deque([(sr, sc, 0)])
            while queue:
                r, c, d = queue.popleft()
                if mat[r][c] == 0:
                    return d
                for nr, nc in ((r + 1, c), (r - 1, c), (r, c + 1), (r, c - 1)):
                    if 0 <= nr < m and 0 <= nc < n and (nr, nc) not in seen:
                        seen.add((nr, nc))
                        queue.append((nr, nc, d + 1))
            return -1

        return [[0 if mat[i][j] == 0 else nearest_zero(i, j) for j in range(n)] for i in range(m)]

Approach 2 · Optimal

Multi-source BFS from all zeros

Time O(m · n)
Space O(m · n)

Seed the queue with every 0 at distance 0 and mark every 1 as unknown. Each BFS step fills in a neighbour one further away.

Set dist = 0 for zeros (enqueue them) and -1 for ones.
Pop a cell; for each neighbour still -1, set dist = current + 1 and enqueue it.

Java

import java.util.*;

class Solution {
    public int[][] updateMatrix(int[][] mat) {
        int m = mat.length;
        int n = mat[0].length;
        int[][] dist = new int[m][n];
        Deque<int[]> queue = new ArrayDeque<>();
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                if (mat[i][j] == 0) {
                    queue.add(new int[] { i, j });
                } else {
                    dist[i][j] = -1;
                }
            }
        }
        int[][] dirs = { { 1, 0 }, { -1, 0 }, { 0, 1 }, { 0, -1 } };
        while (!queue.isEmpty()) {
            int[] cell = queue.poll();
            for (int[] d : dirs) {
                int r = cell[0] + d[0];
                int c = cell[1] + d[1];
                if (r >= 0 && c >= 0 && r < m && c < n && dist[r][c] == -1) {
                    dist[r][c] = dist[cell[0]][cell[1]] + 1;
                    queue.add(new int[] { r, c });
                }
            }
        }
        return dist;
    }
}

Python

from collections import deque

class Solution:
    def updateMatrix(self, mat: List[List[int]]) -> List[List[int]]:
        m, n = len(mat), len(mat[0])
        dist = [[0 if mat[i][j] == 0 else -1 for j in range(n)] for i in range(m)]
        queue = deque((i, j) for i in range(m) for j in range(n) if mat[i][j] == 0)
        while queue:
            r, c = queue.popleft()
            for nr, nc in ((r + 1, c), (r - 1, c), (r, c + 1), (r, c - 1)):
                if 0 <= nr < m and 0 <= nc < n and dist[nr][nc] == -1:
                    dist[nr][nc] = dist[r][c] + 1
                    queue.append((nr, nc))
        return dist

Complexity compared

Approach	Time	Space
Naive BFS from every 1	O((m · n)²)	O(m · n)
Optimal Multi-source BFS from all zeros	O(m · n)	O(m · n)

Key takeaways

"Distance to the nearest X" for every cell: BFS from all X at once.
A two-pass DP (top-left to bottom-right, then back) also works in O(m · n).

Your notes

Saved in this browser as you type.

Solved it on your own?

Mark it done to track your progress, or star it to come back before an interview.

Next problem