rustgym/leetcode/
_1129_shortest_path_with_alternating_colors.rs1struct 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}