1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

//! 木の重心を求める

// 部分木のサイズを求める
pub fn subtree_size(p: usize, u: usize, G: &Vec<Vec<usize>>, W: &[usize], res: &mut [usize]) {
    // 自分を足す
    res[u] += W[u];
    // 集約
    for &v in &G[u] {
        if p == v {
            continue;
        }
        subtree_size(u, v, G, W, res);
        res[u] += res[v];
    }
}

/// 重心を求める
pub fn get_centroid(p: usize, u: usize, G: &Vec<Vec<usize>>, S: &[usize], root: usize) -> usize {
    let mut is_centroid = true;
    // 隣接頂点を調べる
    for &v in &G[u] {
        if v == p {
            continue;
        }
        let res = get_centroid(u, v, G, S, root);
        if res < usize::MAX {
            return res;
        }

        if S[v] > S[root] / 2 {
            is_centroid = false;
        }
    }
    if (S[root] - S[u]) > S[root] / 2 {
        is_centroid = false;
    }

    if is_centroid {
        u
    } else {
        usize::MAX
    }
}