rustgym/leetcode/
_310_minimum_height_trees.rs

1struct Solution;
2use std::collections::VecDeque;
3impl Solution {
4    fn find_min_height_trees(n: i32, edges: Vec<Vec<i32>>) -> Vec<i32> {
5        let n = n as usize;
6        if n == 1 {
7            return vec![0];
8        }
9        let mut graph: Vec<Vec<usize>> = vec![vec![]; n];
10        let mut visited: Vec<bool> = vec![false; n];
11        let mut degree: Vec<usize> = vec![0; n];
12        for e in edges {
13            let u = e[0] as usize;
14            let v = e[1] as usize;
15            graph[u].push(v);
16            graph[v].push(u);
17            degree[u] += 1;
18            degree[v] += 1;
19        }
20
21        let mut leaves: VecDeque<usize> = VecDeque::new();
22        for i in 0..n {
23            if graph[i].len() == 1 {
24                leaves.push_back(i);
25            }
26        }
27        let mut m = n;
28        while m > 2 {
29            m -= leaves.len();
30            for _ in 0..leaves.len() {
31                let u = leaves.pop_front().unwrap();
32                visited[u] = true;
33                for &v in &graph[u] {
34                    if !visited[v] {
35                        degree[v] -= 1;
36                        if degree[v] == 1 {
37                            leaves.push_back(v);
38                        }
39                    }
40                }
41            }
42        }
43        leaves.into_iter().map(|x| x as i32).collect()
44    }
45}
46
47#[test]
48fn test() {
49    let n = 4;
50    let edges = vec_vec_i32![[1, 0], [1, 2], [1, 3]];
51    let mut res = vec![1];
52    let mut ans = Solution::find_min_height_trees(n, edges);
53    ans.sort_unstable();
54    res.sort_unstable();
55    assert_eq!(ans, res);
56    let n = 6;
57    let edges = vec_vec_i32![[0, 3], [1, 3], [2, 3], [4, 3], [5, 4]];
58    let mut res = vec![3, 4];
59    let mut ans = Solution::find_min_height_trees(n, edges);
60    ans.sort_unstable();
61    res.sort_unstable();
62    assert_eq!(ans, res);
63    let n = 3;
64    let edges = vec_vec_i32![[0, 1], [0, 2]];
65    let mut res = vec![0];
66    let mut ans = Solution::find_min_height_trees(n, edges);
67    ans.sort_unstable();
68    res.sort_unstable();
69    assert_eq!(ans, res);
70}