rustgym/leetcode/
_1245_tree_diameter.rs

1struct Solution;
2
3impl Solution {
4    fn tree_diameter(edges: Vec<Vec<i32>>) -> i32 {
5        let n = edges.len() + 1;
6        let mut graph = vec![vec![]; n];
7        for e in edges {
8            let u = e[0] as usize;
9            let v = e[1] as usize;
10            graph[u].push(v);
11            graph[v].push(u);
12        }
13        let mut visited = vec![false; n];
14        let mut res = 0;
15        Self::dfs(0, &mut visited, &mut res, &graph);
16        res as i32
17    }
18
19    fn dfs(u: usize, visited: &mut Vec<bool>, diameter: &mut usize, graph: &[Vec<usize>]) -> usize {
20        visited[u] = true;
21        let mut max_depth = 0;
22        for &v in &graph[u] {
23            if !visited[v] {
24                let depth = Self::dfs(v, visited, diameter, graph);
25                *diameter = (*diameter).max(depth + max_depth);
26                max_depth = max_depth.max(depth);
27            }
28        }
29        max_depth + 1
30    }
31}
32
33#[test]
34fn test() {
35    let edges = vec_vec_i32![[0, 1], [0, 2]];
36    let res = 2;
37    assert_eq!(Solution::tree_diameter(edges), res);
38    let edges = vec_vec_i32![[0, 1], [1, 2], [2, 3], [1, 4], [4, 5]];
39    let res = 4;
40    assert_eq!(Solution::tree_diameter(edges), res);
41}
rustgym/leetcode/_1245_tree_diameter.rs

rustgym/leetcode/
_1245_tree_diameter.rs
