Problem overview
The function compute_grid receives a matrix of 0s and 1s and returns a matrix of the same size where each cell holds the distance to the nearest 0. Moving to an adjacent cell (up, down, left or right) costs 1, and diagonal moves aren't allowed. The matrix contains at least one 0.
For this input:
0 0 0
0 1 0
1 1 1
the output is:
0 0 0
0 1 0
1 2 1
Every 0 stays 0. The center is one step from a zero. The bottom-middle cell is two steps from the nearest zero, because its up, left and right neighbors are all 1s, and the zero above them is two cells away.
Approach: start the search from all zeros at once
The direct idea is to run a search from each 1 until it hits a zero. Every such search can cover much of the matrix, so the total work can grow to roughly the square of the number of cells.
The way out is to turn the question around. Instead of looking for zeros from each one, spread outward from all the zeros together. This is a breadth-first search with many starting points, called multi-source BFS. Put every zero in the queue with distance 0. The cells next to them get distance 1, the cells next to those get distance 2, and so on. Each cell is first reached by the wave that started at its nearest zero, so the distance recorded on first arrival is the final one.
The matrix doubles as the visited record. Before the search, every 1 is changed to -1, which means "no distance yet". Distances are never negative, so there is no confusion. When the search reaches a cell that still holds -1, it writes the distance there and queues the cell.
On the example above, the queue starts with the five zeros. The center cell is reached from a zero and gets 1, and the bottom-left and bottom-right cells get 1 from the zeros above them. The bottom-middle cell is reached from the center (or from a bottom corner) and gets 2.
Step by step
- Store the dimensions:
n = grid.len()rows andm = grid[0].len()columns. Both areusize. - Shadow
gridwith a mutable binding. Parameters are immutable by default, and the distances are written into the grid. - Create the queue
que, an emptyVecDequeof(usize, usize)pairs. - Loop over the rows with
i. - Loop over the columns with
j, inside it. - Queue every zero. If
grid[i][j] == 0, push(i, j)onto the back. Zeros are the sources, and their distance is already right. - Mark every other cell as unvisited by setting it to
-1. - Define
directions, an array of four(i32, i32)offsets: up, down, left and right. They are signed because they can be negative. - Drive the search with
while let, popping the front pair intox1andy1until the queue is empty. Taking from the front is what makes it breadth-first. - Loop over the offsets with
for &(dx, dy) in &directions. - Compute the neighbor
x2andy2ini32, by adding the offsets to the current coordinates cast withas i32. - Check the bounds before using the neighbor: both coordinates must be at least 0 and below
norm(cast toi32for the comparison). - Shadow
x2andy2asusize, now that they are known to be non-negative, since onlyusizecan index. - Look for
-1. Only an unvisited cell is processed, and its first visit is its shortest. - Record the distance as the current cell's plus one and push the neighbor onto the queue.
- Return
gridas the final expression.
The solution in Rust
use std::collections::VecDeque;
impl Solution {
fn compute_grid(grid: Vec<Vec<i32>>) -> Vec<Vec<i32>> {
let n = grid.len();
let m = grid[0].len();
let mut grid = grid;
let mut que: VecDeque<(usize, usize)> = VecDeque::new();
for i in 0..n {
for j in 0..m {
if grid[i][j] == 0 {
que.push_back((i, j));
} else {
grid[i][j] = -1;
}
}
}
let directions: [(i32, i32); 4] = [(-1, 0), (1, 0), (0, -1), (0, 1)];
while let Some((x1, y1)) = que.pop_front() {
for &(dx, dy) in &directions {
let x2 = x1 as i32 + dx;
let y2 = y1 as i32 + dy;
if x2 >= 0 && x2 < n as i32 && y2 >= 0 && y2 < m as i32 {
let (x2, y2) = (x2 as usize, y2 as usize);
if grid[x2][y2] == -1 {
grid[x2][y2] = grid[x1][y1] + 1;
que.push_back((x2, y2));
}
}
}
}
grid
}
}
The Rust details worth a closer look:
- Shadowing.
let mut grid = grid;moves the parameter into a new, mutable binding with the same name. The old immutable one can no longer be used. Writingmut gridin the signature would do the same job. - Why the arithmetic is in
i32. For a cell in row 0,x1 - 1computed as ausizewould subtract from zero. That panics in a debug build and wraps to a huge number in a build with optimizations. Ini32it is simply-1, which the bounds check rejects. - Precedence of
as.x1 as i32 + dxconverts onlyx1, then adds. Likewisex2 < n as i32convertsn. - Tuples. The queue holds
(usize, usize)pairs, andwhile let Some((x1, y1))takes a pair apart as it pops. When the queue is empty,pop_frontreturnsNoneand the loop ends. - Iterating the offsets.
&directionsyields references to the pairs, and the&(dx, dy)pattern copies the two numbers out. - Reading and writing the same grid.
grid[x2][y2] = grid[x1][y1] + 1;compiles because the right side is evaluated to a plaini32first.
Complexity
Time: O(n · m). Setting up visits every cell once. In the search each cell enters the queue at most once, because it is no longer -1 afterwards, and each dequeued cell checks four neighbors.
Space: O(n · m). The queue can hold many cells at once, for example all the zeros at the start. Reusing the input as the answer saves a second matrix.
Edge cases and common mistakes
- All zeros. The output equals the input, and no neighbor is ever
-1. - A single zero in a corner. For
[[1,1,1],[1,1,1],[1,1,0]]the answer is[[4,3,2],[3,2,1],[2,1,0]]: the distances grow with the sum of the steps down and right. - One row or one column.
[[0,1,1,1,1]]gives[[0,1,2,3,4]], and[[1],[1],[0]]gives[[2],[1],[0]]. The bounds check handles the missing neighbors. - Using a stack instead of a queue. Popping from the back makes it a depth-first search, and cells can get distances that are too large.
- Leaving the ones as 1. The check for unvisited cells would then match nothing useful, and a real distance of 1 couldn't be told apart from "not visited yet". The
-1marker is what keeps them apart. - Starting the search from the ones. This repeats work for every
1and is much slower on large matrices.
Testing
[[0,0,0],[0,1,0],[0,0,0]]gives[[0,0,0],[0,1,0],[0,0,0]].[[0,0,0],[0,1,0],[1,1,1]]gives[[0,0,0],[0,1,0],[1,2,1]].[[1,1,1],[1,1,1],[1,1,0]]gives[[4,3,2],[3,2,1],[2,1,0]].[[0,1,1,1,1]]gives[[0,1,2,3,4]].
