library_checker/tree/
tree_diameter.rs

1use competitive::prelude::*;
2use competitive::tree::XorLinkedRootedTree;
3
4#[verify::library_checker("tree_diameter")]
5pub fn tree_diameter(reader: impl Read, writer: impl Write) {
6    prepare_io!(reader, writer);
7    sc!(n, edges: [(u32, u32, u64); n - 1]);
8    let mut parent = vec![n as u32; n];
9    let mut farthest: Vec<_> = (0..n).map(|u| (0, u)).collect();
10    let (mut diameter, mut left, mut right, mut center) = (0, 0, 0, 0);
11    XorLinkedRootedTree::builder(n).with_eindexed().run(
12        0,
13        edges.iter().map(|&(u, v, _)| (u as usize, v as usize)),
14        |u, p, e| {
15            parent[u] = p as u32;
16            let distance = farthest[u].0 + edges[e].2;
17            if diameter < distance + farthest[p].0 {
18                diameter = distance + farthest[p].0;
19                left = farthest[u].1;
20                right = farthest[p].1;
21                center = p;
22            }
23            if farthest[p].0 < distance {
24                farthest[p] = (distance, farthest[u].1);
25            }
26        },
27    );
28    let mut path = Vec::new();
29    while left != center {
30        path.push(left);
31        left = parent[left] as usize;
32    }
33    path.push(center);
34    let middle = path.len();
35    while right != center {
36        path.push(right);
37        right = parent[right] as usize;
38    }
39    path[middle..].reverse();
40    pp!(diameter, path.len(); @it path);
41}
library_checker/tree/tree_diameter.rs

library_checker/tree/
tree_diameter.rs
