cp_library_rs/graph/
centroid.rs

1//! 木の重心を求める
2
3// 部分木のサイズを求める
4pub fn subtree_size(p: usize, u: usize, G: &Vec<Vec<usize>>, W: &[usize], res: &mut [usize]) {
5    // 自分を足す
6    res[u] += W[u];
7    // 集約
8    for &v in &G[u] {
9        if p == v {
10            continue;
11        }
12        subtree_size(u, v, G, W, res);
13        res[u] += res[v];
14    }
15}
16
17/// 重心を求める
18pub fn get_centroid(p: usize, u: usize, G: &Vec<Vec<usize>>, S: &[usize], root: usize) -> usize {
19    let mut is_centroid = true;
20    // 隣接頂点を調べる
21    for &v in &G[u] {
22        if v == p {
23            continue;
24        }
25        let res = get_centroid(u, v, G, S, root);
26        if res < usize::MAX {
27            return res;
28        }
29
30        if S[v] > S[root] / 2 {
31            is_centroid = false;
32        }
33    }
34    if (S[root] - S[u]) > S[root] / 2 {
35        is_centroid = false;
36    }
37
38    if is_centroid {
39        u
40    } else {
41        usize::MAX
42    }
43}