01 Matrix LeetCode solution in Rust

1
use std::collections::VecDeque;
2

3
impl Solution {
4
  pub fn update_matrix(mat: Vec<Vec<i32>>) -> Vec<Vec<i32>> {
5
    let (r, c) = (mat.len(), mat[0].len());
6
    let mut answer = vec![vec![i32::MAX; c]; r];
7
    let mut vd = VecDeque::new();
8
    for (i, row) in mat.iter().enumerate() {
9
      for (j, col) in row.iter().enumerate() {
10
        if *col == 0 {
11
          answer[i][j] = 0;
12
          vd.push_back((i, j));
13
        }
14
      }
15
    }
16
    while let Some((i, j)) = vd.pop_front() {
17
      for d in [0, 1, 0, !0, 0].windows(2) {
18
        let di = i.wrapping_add(d[0]);
19
        let dj = j.wrapping_add(d[1]);
20
        if di < r && dj < c && answer[di][dj] > answer[i][j] {
21
          answer[di][dj] = answer[i][j] + 1;
22
          vd.push_back((di, dj));
23
        }
24
      }
25
    }
26
    answer
27
  }
28
}