rustgym/leetcode/
_1129_shortest_path_with_alternating_colors.rs

1struct Solution;
2use std::collections::VecDeque;
3
4impl Solution {
5    fn shortest_alternating_paths(
6        n: i32,
7        red_edges: Vec<Vec<i32>>,
8        blue_edges: Vec<Vec<i32>>,
9    ) -> Vec<i32> {
10        let n = n as usize;
11        let mut graph: Vec<Vec<Vec<usize>>> = vec![vec![vec![]; n]; 2];
12        let mut visited: Vec<Vec<bool>> = vec![vec![false; n]; 2];
13        let mut res: Vec<i32> = vec![-1; n];
14        for e in red_edges {
15            let u = e[0] as usize;
16            let v = e[1] as usize;
17            graph[0][u].push(v);
18        }
19        for e in blue_edges {
20            let u = e[0] as usize;
21            let v = e[1] as usize;
22            graph[1][u].push(v);
23        }
24        let mut queue: VecDeque<(usize, usize, usize)> = VecDeque::new();
25        queue.push_back((0, 0, 0));
26        queue.push_back((0, 1, 0));
27        visited[0][0] = true;
28        visited[1][0] = true;
29        while let Some((u, c, d)) = queue.pop_front() {
30            if res[u] == -1 {
31                res[u] = d as i32;
32            }
33            for &v in &graph[c][u] {
34                if !visited[1 - c][v] {
35                    visited[1 - c][v] = true;
36                    queue.push_back((v, 1 - c, d + 1));
37                }
38            }
39        }
40        res
41    }
42}
43
44#[test]
45fn test() {
46    let n = 3;
47    let red_edges = vec_vec_i32![[0, 1], [1, 2]];
48    let blue_edges = vec_vec_i32![];
49    let res = vec![0, 1, -1];
50    assert_eq!(
51        Solution::shortest_alternating_paths(n, red_edges, blue_edges),
52        res
53    );
54    let n = 3;
55    let red_edges = vec_vec_i32![[0, 1]];
56    let blue_edges = vec_vec_i32![[2, 1]];
57    let res = vec![0, 1, -1];
58    assert_eq!(
59        Solution::shortest_alternating_paths(n, red_edges, blue_edges),
60        res
61    );
62    let n = 3;
63    let red_edges = vec_vec_i32![[1, 0]];
64    let blue_edges = vec_vec_i32![[2, 1]];
65    let res = vec![0, -1, -1];
66    assert_eq!(
67        Solution::shortest_alternating_paths(n, red_edges, blue_edges),
68        res
69    );
70    let n = 3;
71    let red_edges = vec_vec_i32![[0, 1]];
72    let blue_edges = vec_vec_i32![[1, 2]];
73    let res = vec![0, 1, 2];
74    assert_eq!(
75        Solution::shortest_alternating_paths(n, red_edges, blue_edges),
76        res
77    );
78    let n = 3;
79    let red_edges = vec_vec_i32![[0, 1], [0, 2]];
80    let blue_edges = vec_vec_i32![[1, 0]];
81    let res = vec![0, 1, 1];
82    assert_eq!(
83        Solution::shortest_alternating_paths(n, red_edges, blue_edges),
84        res
85    );
86}