tree_diameter/
lib.rs

1use graph::is_tree;
2
3/// 木の直径と直径をなす頂点のベクタを返します。
4pub fn tree_diameter(n: usize, edges: &[(usize, usize, u64)]) -> (u64, Vec<usize>) {
5    if n == 0 {
6        return (0, Vec::new());
7    }
8
9    assert!(is_tree(
10        n,
11        &edges
12            .iter()
13            .copied()
14            .map(|(u, v, _)| (u, v))
15            .collect::<Vec<_>>()
16    ));
17
18    for &(_, _, c) in edges {
19        assert!(c >= 1);
20    }
21
22    let mut graph = vec![vec![]; n];
23    for &(u, v, c) in edges {
24        graph[u].push((v, c));
25        graph[v].push((u, c));
26    }
27
28    fn dfs(
29        i: usize,
30        p: usize,
31        g: &[Vec<(usize, u64)>],
32        dist: &mut Vec<u64>,
33        parent: &mut Vec<usize>,
34    ) {
35        parent[i] = p;
36        for &(j, c) in &g[i] {
37            if j == p {
38                continue;
39            }
40            dist[j] = dist[i] + c;
41            dfs(j, i, g, dist, parent);
42        }
43    }
44
45    let mut dist = vec![0; n];
46    let mut parent = vec![0; n];
47    dfs(0, 0, &graph, &mut dist, &mut parent);
48
49    let max_dist = dist.iter().max().copied().unwrap();
50    let s = (0..n).position(|i| dist[i] == max_dist).unwrap();
51    dist = vec![0; n];
52    parent = vec![s; n];
53    dfs(s, s, &graph, &mut dist, &mut parent);
54
55    let diameter = dist.iter().max().copied().unwrap();
56    let t = (0..n).position(|i| dist[i] == diameter).unwrap();
57
58    let mut cur = t;
59    let mut path = Vec::new();
60    path.push(cur);
61    while cur != parent[cur] {
62        cur = parent[cur];
63        path.push(cur);
64    }
65    (diameter, path)
66}
67
68#[cfg(test)]
69mod tests {
70    use super::tree_diameter;
71
72    #[test]
73    fn test_small() {
74        assert_eq!(tree_diameter(0, &[]), (0, vec![]));
75        assert_eq!(tree_diameter(1, &[]), (0, vec![0]));
76
77        let (diameter, path) = tree_diameter(2, &[(0, 1, 1)]);
78        assert_eq!(diameter, 1);
79        assert!(path == vec![0, 1] || path == vec![1, 0]);
80
81        let (diameter, path) = tree_diameter(3, &[(0, 1, 100), (0, 2, 100)]);
82        assert_eq!(diameter, 200);
83        assert!(path == vec![1, 0, 2] || path == vec![2, 0, 1]);
84    }
85}
tree_diameter/lib.rs

tree_diameter/
lib.rs
