haar_lib/tree/
utils.rs

1//! 木上の便利関数
2use crate::num::one_zero::Zero;
3use crate::tree::*;
4use std::ops::Add;
5
6/// rootを根としたときの根から各頂点への距離を列挙する。
7///
8/// **Time complexity** $O(n)$
9pub fn tree_distance<E: TreeEdgeTrait>(tr: &Tree<E>, root: usize) -> Vec<E::Weight>
10where
11    E::Weight: Add<Output = E::Weight> + Copy + Zero,
12{
13    let n = tr.len();
14    let mut ret = vec![E::Weight::zero(); n];
15    let mut check = vec![false; n];
16    let mut stack = vec![root];
17
18    while let Some(cur) = stack.pop() {
19        check[cur] = true;
20
21        tr.nodes[cur]
22            .neighbors()
23            .filter(|e| !check[e.to()])
24            .for_each(|e| {
25                ret[e.to()] = ret[cur] + e.weight();
26                stack.push(e.to());
27            });
28    }
29
30    ret
31}
32
33/// 木の任意の2頂点の距離の最大値を求める。
34///
35/// **Time complexity** $O(n)$
36pub fn tree_diameter<E: TreeEdgeTrait>(tr: &Tree<E>) -> (E::Weight, usize, usize)
37where
38    E::Weight: Add<Output = E::Weight> + Copy + Zero + Ord,
39{
40    let a = tree_distance(tr, 0);
41    let (u, _) = a
42        .iter()
43        .enumerate()
44        .max_by(|(_, x), (_, y)| x.cmp(y))
45        .unwrap();
46
47    let b = tree_distance(tr, u);
48    let (v, &d) = b
49        .iter()
50        .enumerate()
51        .max_by(|(_, x), (_, y)| x.cmp(y))
52        .unwrap();
53
54    (d, u, v)
55}
56
57/// 木の各頂点について、そこからの距離の最大値を列挙する。
58///
59/// **Time complexity** $O(n)$
60pub fn tree_height<E: TreeEdgeTrait>(tr: &Tree<E>) -> Vec<(E::Weight, usize)>
61where
62    E::Weight: Add<Output = E::Weight> + Copy + Zero + Ord,
63{
64    let d = tree_distance(tr, 0);
65    let (u, _) = d
66        .iter()
67        .enumerate()
68        .max_by(|(_, x), (_, y)| x.cmp(y))
69        .unwrap();
70    let d1 = tree_distance(tr, u);
71    let (v, _) = d1
72        .iter()
73        .enumerate()
74        .max_by(|(_, x), (_, y)| x.cmp(y))
75        .unwrap();
76    let d2 = tree_distance(tr, v);
77
78    d1.into_iter()
79        .zip(d2)
80        .map(|(x, y)| if x > y { (x, u) } else { (y, v) })
81        .collect()
82}
83
84/// 木上の2頂点を結ぶパス上の頂点列を求める。
85///
86/// **Time complexity** $O(n)$
87pub fn tree_path<E: TreeEdgeTrait>(tr: &Tree<E>, u: usize, v: usize) -> Vec<usize> {
88    let n = tr.len();
89    let mut ret = vec![];
90    let mut stack = vec![];
91    let mut check = vec![false; n];
92
93    stack.push((u, 0));
94
95    while let Some((i, st)) = stack.pop() {
96        match st {
97            0 => {
98                stack.push((i, 1));
99                ret.push(i);
100
101                if i == v {
102                    break;
103                }
104
105                check[i] = true;
106
107                tr.nodes[i]
108                    .neighbors()
109                    .filter(|e| !check[e.to()])
110                    .for_each(|e| stack.push((e.to(), 0)));
111            }
112            1 => {
113                ret.pop();
114            }
115            _ => unreachable!(),
116        }
117    }
118
119    ret
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125
126    #[test]
127    fn test_diameter() {
128        // https://onlinejudge.u-aizu.ac.jp/courses/library/5/GRL/all/GRL_5_A
129        let mut builder = TreeBuilder::new(4);
130        builder.extend(
131            vec![(0, 1, 2), (1, 2, 1), (1, 3, 3)]
132                .into_iter()
133                .map(|(u, v, w)| TreeEdge::new(u, v, w, ())),
134        );
135        let tree = builder.build();
136        assert_eq!(tree_diameter(&tree).0, 5);
137
138        let mut builder = TreeBuilder::new(4);
139        builder.extend(
140            vec![(0, 1, 1), (1, 2, 2), (2, 3, 4)]
141                .into_iter()
142                .map(|(u, v, w)| TreeEdge::new(u, v, w, ())),
143        );
144        let tree = builder.build();
145        assert_eq!(tree_diameter(&tree).0, 7);
146    }
147
148    #[test]
149    fn test_height() {
150        // https://onlinejudge.u-aizu.ac.jp/courses/library/5/GRL/all/GRL_5_B
151        let mut builder = TreeBuilder::new(4);
152        builder.extend(
153            vec![(0, 1, 2), (1, 2, 1), (1, 3, 3)]
154                .into_iter()
155                .map(|(u, v, w)| TreeEdge::new(u, v, w, ())),
156        );
157        let tree = builder.build();
158        assert_eq!(
159            tree_height(&tree)
160                .into_iter()
161                .map(|(x, _)| x)
162                .collect::<Vec<_>>(),
163            [5, 3, 4, 5]
164        );
165    }
166}
haar_lib/tree/utils.rs

haar_lib/tree/
utils.rs
