competitive_library/graph/
tree_diameter.rs

1pub trait Edge {
2    type T: Clone;
3    fn get_edge(a: &Self::T) -> usize;
4    fn get_cost(a: &Self::T) -> i64;
5}
6
7pub struct UnWeightedEdge {}
8impl Edge for UnWeightedEdge {
9    type T = usize;
10    #[inline]
11    fn get_edge(a: &Self::T) -> usize {
12        *a
13    }
14    #[inline]
15    fn get_cost(_: &Self::T) -> i64 {
16        1
17    }
18}
19pub struct WeightedEdge {}
20impl Edge for WeightedEdge {
21    type T = (usize, i64);
22
23    #[inline]
24    fn get_edge(a: &Self::T) -> usize {
25        a.0
26    }
27    #[inline]
28    fn get_cost(a: &Self::T) -> i64 {
29        a.1
30    }
31}
32
33pub fn tree_diameter<E: Edge>(e: &[Vec<E::T>]) -> (i64, Vec<usize>) {
34    let (_, path_1) = bfs::<E>(e, 0);
35    bfs::<E>(e, path_1[path_1.len() - 1])
36}
37
38#[deprecated(since = "0.1.0", note = "use `tree_diameter` instead")]
39pub use tree_diameter as tree_diamiter;
40
41#[inline]
42fn bfs<E: Edge>(e: &[Vec<E::T>], start: usize) -> (i64, Vec<usize>) {
43    let mut que = std::collections::VecDeque::new();
44    let mut max_cost_index = (start, 0);
45    let mut previous = vec![None; e.len()];
46    previous[start] = Some(start);
47
48    que.push_back((start, 0));
49
50    while let Some((v, cost)) = que.pop_front() {
51        for edge in e[v].iter() {
52            let to = E::get_edge(edge);
53            if previous[to].is_some() {
54                continue;
55            }
56            previous[to] = Some(v);
57            que.push_back((E::get_edge(edge), cost + E::get_cost(edge)));
58        }
59
60        if max_cost_index.1 < cost {
61            max_cost_index = (v, cost);
62        }
63    }
64    (max_cost_index.1, restore_path(max_cost_index.0, &previous))
65}
66
67#[inline]
68fn restore_path(end: usize, previous: &[Option<usize>]) -> Vec<usize> {
69    let mut buff = end;
70    let mut v = vec![buff];
71
72    while let Some(i) = previous[buff] {
73        if buff == i {
74            break;
75        }
76        buff = i;
77        v.push(buff);
78    }
79    v.reverse();
80    v
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86    #[test]
87    fn test_tree_diameter() {
88        let n = 8;
89
90        let input = vec![
91            (0, 1, 5),
92            (1, 2, 3),
93            (2, 3, 1),
94            (1, 4, 2),
95            (4, 7, 4),
96            (1, 5, 7),
97            (2, 6, 5),
98        ];
99        let mut e = vec![vec![]; n];
100        for (a, b, c) in input {
101            e[a].push((b, c));
102            e[b].push((a, c));
103        }
104
105        let a = tree_diameter::<WeightedEdge>(&e);
106        assert_eq!(&a.1, &[6, 2, 1, 5]);
107        assert_eq!(a.0, 15);
108    }
109}
competitive_library/graph/tree_diameter.rs

competitive_library/graph/
tree_diameter.rs
