competitive/tree/
centroid_decomposition.rs

1use super::{ConvolveSteps, U64Convolve, UndirectedSparseGraph};
2use std::mem::swap;
3
4fn centroid_decomposition_dfs(
5    parents: &[usize],
6    vs: &[usize],
7    f: &mut impl FnMut(&[usize], &[usize], usize, usize),
8) {
9    let n = vs.len();
10    assert!(n >= 2);
11    if n == 2 {
12        return;
13    }
14    let mut size = vec![1; n];
15    let mut c = !0;
16    for i in (0..n).rev() {
17        if size[i] >= n.div_ceil(2) {
18            c = i;
19            break;
20        }
21        size[parents[i]] += size[i];
22    }
23    let mut color = vec![!0u8; n];
24    let mut order = vec![!0usize; n];
25    order[c] = 0;
26    let mut count = 1usize;
27    let mut sum_size = 0usize;
28    for u in 1..n {
29        if parents[u] == c && sum_size + size[u] <= (n - 1) / 2 {
30            sum_size += size[u];
31            color[u] = 0;
32            order[u] = count;
33            count += 1;
34        }
35    }
36    for u in 1..n {
37        if color[parents[u]] == 0 {
38            color[u] = 0;
39            order[u] = count;
40            count += 1;
41        }
42    }
43    let lsize = count - 1;
44    {
45        let mut u = parents[c];
46        while u != !0 {
47            color[u] = 1;
48            order[u] = count;
49            count += 1;
50            u = parents[u];
51        }
52    }
53    for u in 0..n {
54        if u != c && color[u] == !0 {
55            color[u] = 1;
56            order[u] = count;
57            count += 1;
58        }
59    }
60    assert_eq!(count, n);
61    let rsize = n - lsize - 1;
62    let mut cparents = vec![!0usize; n];
63    let mut cvs = vec![!0usize; n];
64    let mut lparents = vec![!0usize; lsize + 1];
65    let mut lvs = vec![!0usize; lsize + 1];
66    let mut rparents = vec![!0usize; rsize + 1];
67    let mut rvs = vec![!0usize; rsize + 1];
68    for u in 0..n {
69        let i = order[u];
70        cvs[i] = vs[u];
71        if color[u] != 1 {
72            lvs[i] = vs[u];
73        }
74        if color[u] != 0 {
75            rvs[if i == 0 { 0 } else { i - lsize }] = vs[u];
76        }
77    }
78    for u in 1..n {
79        let mut x = order[u];
80        let mut y = order[parents[u]];
81        if x > y {
82            swap(&mut x, &mut y);
83        }
84        cparents[y] = x;
85        if color[u] != 1 && color[parents[u]] != 1 {
86            lparents[y] = x;
87        }
88        if color[u] != 0 && color[parents[u]] != 0 {
89            rparents[if y == 0 { 0 } else { y - lsize }] = if x == 0 { 0 } else { x - lsize };
90        }
91    }
92    f(&cparents, &cvs, lsize, rsize);
93    centroid_decomposition_dfs(&lparents, &lvs, f);
94    centroid_decomposition_dfs(&rparents, &rvs, f);
95}
96
97impl UndirectedSparseGraph {
98    /// 1/3 centroid decomposition
99    ///
100    /// - f: (parents: &[usize], vs: &[usize], lsize: usize, rsize: usize)
101    /// - 0: root, 1..=lsize: left subtree, lsize+1..=lsize+rsize: right subtree
102    pub fn centroid_decomposition(&self, mut f: impl FnMut(&[usize], &[usize], usize, usize)) {
103        let n = self.vertices_size();
104        if n <= 1 {
105            return;
106        }
107        let mut vs = Vec::with_capacity(n);
108        let mut parents = vec![!0; n];
109        vs.push(0usize);
110        for l in 0..n {
111            let u = vs[l];
112            for a in self.adjacencies(u) {
113                if a.to != parents[u] {
114                    vs.push(a.to);
115                    parents[a.to] = u;
116                }
117            }
118        }
119        let mut new_idx = vec![0; n];
120        for i in 0..n {
121            new_idx[vs[i]] = i;
122        }
123        let mut new_parent = vec![!0; n];
124        for i in 1..n {
125            new_parent[new_idx[i]] = new_idx[parents[i]];
126        }
127        centroid_decomposition_dfs(&new_parent, &vs, &mut f);
128    }
129
130    pub fn distance_frequencies(&self) -> Vec<u64> {
131        let n = self.vertices_size();
132        let mut table = vec![0u64; n];
133        if n == 0 {
134            return table;
135        }
136        table[0] = n as u64;
137        if n == 1 {
138            return table;
139        }
140        table[1] = (n * 2 - 2) as u64;
141        self.centroid_decomposition(|parents, vs, lsize, _rsize| {
142            let n = vs.len();
143            let mut dist = vec![0usize; n];
144            for i in 1..n {
145                dist[i] = dist[parents[i]] + 1;
146            }
147            let d_max = dist.iter().max().cloned().unwrap_or_default();
148            let mut f = vec![0u64; d_max + 1];
149            let mut g = vec![0u64; d_max + 1];
150            for i in 1..=lsize {
151                f[dist[i]] += 1;
152            }
153            for i in lsize + 1..n {
154                g[dist[i]] += 1;
155            }
156            while f.last().is_some_and(|&x| x == 0) {
157                f.pop();
158            }
159            while g.last().is_some_and(|&x| x == 0) {
160                g.pop();
161            }
162            let h = U64Convolve::convolve(f, g);
163            for (i, &x) in h.iter().enumerate() {
164                table[i] += x * 2;
165            }
166        });
167        table
168    }
169}
170
171#[cfg(test)]
172mod tests {
173    use crate::{tools::Xorshift, tree::MixedTree};
174
175    #[test]
176    fn test_distance_frequencies() {
177        let mut rng = Xorshift::default();
178        for _ in 0..200 {
179            let g = rng.random(MixedTree(1usize..100));
180            let n = g.vertices_size();
181            let result = g.distance_frequencies();
182            let mut expected = vec![0u64; n];
183            for u in 0..n {
184                let depth = g.tree_depth(u);
185                for v in 0..n {
186                    expected[depth[v] as usize] += 1;
187                }
188            }
189            assert_eq!(result, expected);
190        }
191    }
192}
competitive/tree/centroid_decomposition.rs

competitive/tree/
centroid_decomposition.rs
