haar_lib/tree/
centroid_decomposition.rs

1//! 重心分解
2//!
3//! # Problems
4//! - <https://atcoder.jp/contests/abc291/tasks/abc291_h>
5//! - <https://judge.yosupo.jp/problem/frequency_table_of_tree_distance>
6use crate::tree::*;
7
8/// [`CentroidDecomposition`]の頂点ノード
9#[derive(Clone)]
10pub struct Node {
11    /// 親の頂点
12    pub par: Option<usize>,
13    /// 子の頂点列
14    pub children: Vec<usize>,
15    /// 深さ
16    pub depth: usize,
17    /// 部分木の大きさ
18    pub subsize: usize,
19}
20
21/// 重心分解
22pub struct CentroidDecomposition {
23    nodes: Vec<Node>,
24}
25
26impl CentroidDecomposition {
27    /// 木`tree`を重心分解する。
28    pub fn new<E: TreeEdgeTrait>(tree: &Tree<E>) -> Self {
29        let n = tree.len();
30        let mut this = Self {
31            nodes: vec![
32                Node {
33                    par: None,
34                    children: vec![],
35                    depth: !0,
36                    subsize: 0,
37                };
38                n
39            ],
40        };
41
42        let mut subsize = vec![0; n];
43        this.decompose(tree, 0, None, 0, &mut subsize);
44        for (a, s) in this.nodes.iter_mut().zip(subsize) {
45            a.subsize = s;
46        }
47        this
48    }
49
50    /// 重心分解後の頂点列への参照を返す。
51    pub fn nodes(&self) -> &[Node] {
52        &self.nodes
53    }
54
55    fn decompose<E: TreeEdgeTrait>(
56        &mut self,
57        tree: &Tree<E>,
58        cur: usize,
59        par: Option<usize>,
60        d: usize,
61        subsize: &mut [usize],
62    ) {
63        self.dfs_subsize(tree, cur, None, subsize);
64        let c = self.get_centroid(tree, cur, None, subsize[cur], subsize);
65
66        self.nodes[c].par = par;
67        self.nodes[c].depth = d;
68
69        if let Some(par) = par {
70            self.nodes[par].children.push(c);
71        }
72
73        for e in tree.nodes[c].neighbors() {
74            if self.nodes[e.to()].depth != !0 {
75                continue;
76            }
77            self.decompose(tree, e.to(), Some(c), d + 1, subsize);
78        }
79    }
80
81    fn get_centroid<E: TreeEdgeTrait>(
82        &self,
83        tree: &Tree<E>,
84        cur: usize,
85        par: Option<usize>,
86        total_size: usize,
87        subsize: &[usize],
88    ) -> usize {
89        tree.nodes[cur]
90            .neighbors()
91            .filter(|e| par != Some(e.to()) && self.nodes[e.to()].depth == !0)
92            .find(|e| 2 * subsize[e.to()] > total_size)
93            .map_or(cur, |e| {
94                self.get_centroid(tree, e.to(), Some(cur), total_size, subsize)
95            })
96    }
97
98    fn dfs_subsize<E: TreeEdgeTrait>(
99        &self,
100        tree: &Tree<E>,
101        cur: usize,
102        par: Option<usize>,
103        subsize: &mut [usize],
104    ) {
105        subsize[cur] = 1 + tree.nodes[cur]
106            .neighbors()
107            .filter(|e| par != Some(e.to()) && self.nodes[e.to()].depth == !0)
108            .map(|e| {
109                self.dfs_subsize(tree, e.to(), Some(cur), subsize);
110                subsize[e.to()]
111            })
112            .sum::<usize>();
113    }
114}
haar_lib/tree/centroid_decomposition.rs

haar_lib/tree/
centroid_decomposition.rs
