struct Solution;
use std::collections::VecDeque;

impl Solution {
    fn update_matrix(mut matrix: Vec<Vec<i32>>) -> Vec<Vec<i32>> {
        let n = matrix.len();
        let m = matrix[0].len();
        let mut visited: Vec<Vec<bool>> = vec![vec![false; m]; n];
        let mut queue: VecDeque<(usize, usize, i32)> = VecDeque::new();
        for i in 0..n {
            for j in 0..m {
                if matrix[i][j] == 0 {
                    queue.push_back((i, j, 0));
                }
            }
        }
        while let Some((i, j, d)) = queue.pop_front() {
            if visited[i][j] {
                continue;
            }
            visited[i][j] = true;
            matrix[i][j] = d;
            if i > 0 {
                queue.push_back((i - 1, j, d + 1));
            }
            if j > 0 {
                queue.push_back((i, j - 1, d + 1));
            }
            if i + 1 < n {
                queue.push_back((i + 1, j, d + 1));
            }
            if j + 1 < m {
                queue.push_back((i, j + 1, d + 1));
            }
        }
        matrix
    }
}

#[test]
fn test() {
    let matrix = vec_vec_i32![[0, 0, 0], [0, 1, 0], [0, 0, 0]];
    let res = vec_vec_i32![[0, 0, 0], [0, 1, 0], [0, 0, 0]];
    assert_eq!(Solution::update_matrix(matrix), res);
    let matrix = vec_vec_i32![[0, 0, 0], [0, 1, 0], [1, 1, 1]];
    let res = vec_vec_i32![[0, 0, 0], [0, 1, 0], [1, 2, 1]];
    assert_eq!(Solution::update_matrix(matrix), res);
    let matrix = vec_vec_i32![[0], [0], [0], [0], [0]];
    let res = vec_vec_i32![[0], [0], [0], [0], [0]];
    assert_eq!(Solution::update_matrix(matrix), res);
}