haar_lib/tree/
centroid.rs

1//! 重心列挙
2
3use crate::tree::*;
4/// 木の重心を列挙する
5pub fn centroids<E: TreeEdgeTrait>(tree: &Tree<E>) -> Vec<usize> {
6    let n = tree.len();
7    let mut sub = vec![0; n];
8    let mut ret = vec![];
9    dfs(tree, &mut sub, &mut ret, n, 0, None);
10    ret
11}
12
13fn dfs<E: TreeEdgeTrait>(
14    tree: &Tree<E>,
15    sub: &mut [usize],
16    ret: &mut Vec<usize>,
17    size: usize,
18    cur: usize,
19    par: Option<usize>,
20) {
21    sub[cur] = 1;
22
23    let mut check = true;
24
25    for e in tree.nodes[cur]
26        .neighbors()
27        .filter(|e| par.is_none_or(|p| p != e.to()))
28    {
29        dfs(tree, sub, ret, size, e.to(), Some(cur));
30
31        if sub[e.to()] > size / 2 {
32            check = false;
33        }
34        sub[cur] += sub[e.to()];
35    }
36
37    if size - sub[cur] > size / 2 {
38        check = false;
39    }
40
41    if check {
42        ret.push(cur);
43    }
44}
45
46#[cfg(test)]
47mod tests {
48    use super::*;
49
50    #[test]
51    fn test() {
52        let mut builder = TreeBuilder::new(3);
53        builder.extend(
54            vec![(0, 1), (1, 2)]
55                .into_iter()
56                .map(|(u, v)| TreeEdge::new(u, v, (), ())),
57        );
58        let tree = builder.build();
59        assert_eq!(centroids(&tree), vec![1]);
60
61        let mut builder = TreeBuilder::new(4);
62        builder.extend(
63            vec![(0, 1), (1, 2), (2, 3)]
64                .into_iter()
65                .map(|(u, v)| TreeEdge::new(u, v, (), ())),
66        );
67        let tree = builder.build();
68        let mut ans = centroids(&tree);
69        ans.sort();
70        assert_eq!(ans, vec![1, 2]);
71    }
72}
haar_lib/tree/centroid.rs

haar_lib/tree/
centroid.rs
