competitive/tree/
distance_frequencies.rs

1use super::{ConvolveSteps, U64Convolve, UndirectedSparseGraph};
2
3impl UndirectedSparseGraph {
4    pub fn distance_frequencies(&self) -> Vec<u64> {
5        let n = self.vertices_size();
6        let mut table = vec![0u64; n];
7        if n == 0 {
8            return table;
9        }
10        table[0] = n as u64;
11        if n == 1 {
12            return table;
13        }
14        table[1] = (n * 2 - 2) as u64;
15        self.centroid_decomposition(|parents, vs, lsize, _rsize| {
16            let n = vs.len();
17            let mut dist = vec![0usize; n];
18            for i in 1..n {
19                dist[i] = dist[parents[i]] + 1;
20            }
21            let d_max = dist.iter().max().cloned().unwrap_or_default();
22            let mut f = vec![0u64; d_max + 1];
23            let mut g = vec![0u64; d_max + 1];
24            for i in 1..=lsize {
25                f[dist[i]] += 1;
26            }
27            for i in lsize + 1..n {
28                g[dist[i]] += 1;
29            }
30            while f.last().is_some_and(|&x| x == 0) {
31                f.pop();
32            }
33            while g.last().is_some_and(|&x| x == 0) {
34                g.pop();
35            }
36            let h = U64Convolve::convolve(f, g);
37            for (i, &x) in h.iter().enumerate() {
38                table[i] += x * 2;
39            }
40        });
41        table
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use crate::{tools::Xorshift, tree::MixedTree};
48    #[test]
49    fn test_distance_frequencies() {
50        let mut rng = Xorshift::default();
51        for _ in 0..200 {
52            let g = rng.random(MixedTree(1usize..100));
53            let n = g.vertices_size();
54            let result = g.distance_frequencies();
55            let mut expected = vec![0u64; n];
56            for u in 0..n {
57                let depth = g.tree_depth(u);
58                for v in 0..n {
59                    expected[depth[v] as usize] += 1;
60                }
61            }
62            assert_eq!(result, expected);
63        }
64    }
65}
competitive/tree/distance_frequencies.rs

competitive/tree/
distance_frequencies.rs
