add cent decomp (#38)

* add cent decomp

* complete docs

* fix docs

* increase TL for naive conv

* clean up dfs order to a bfs

* single core

* undo ci changes

* add ripped fft

* fix clippy

* fix doc

Author: cameroncuster

Cargo.toml

@@ -46,6 +46,10 @@ path = "examples/data_structures/seg_tree_build_on_array.rs"
 name = "trie"
 path = "examples/data_structures/trie.rs"
 
+[[example]]
+name = "count_paths_per_length"
+path = "examples/graphs/count_paths_per_length.rs"
+
 [[example]]
 name = "dijk_aizu"
 path = "examples/graphs/dijk_aizu.rs"

examples/graphs/count_paths_per_length.rs

@@ -0,0 +1,32 @@
+// verification-helper: PROBLEM https://judge.yosupo.jp/problem/frequency_table_of_tree_distance
+
+use proconio::input;
+use programming_team_code_rust::graphs::count_paths_per_length::count_paths_per_length;
+
+fn main() {
+    input! {
+        n: usize
+    }
+
+    let mut adj = vec![vec![]; n];
+    for _ in 0..n - 1 {
+        input! {
+            u: usize,
+            v: usize
+        }
+        adj[u].push(v);
+        adj[v].push(u);
+    }
+
+    let paths_per_length = count_paths_per_length(&adj);
+
+    println!(
+        "{}",
+        paths_per_length
+            .iter()
+            .skip(1)
+            .map(|&x| x.to_string())
+            .collect::<Vec<_>>()
+            .join(" ")
+    );
+}

src/graphs/cent_decomp.rs

@@ -0,0 +1,66 @@
+//! # Centroid Decomposition
+
+fn calc_sz(adj: &[Vec<usize>], u: usize, p: usize, sub_sz: &mut [usize]) {
+    sub_sz[u] = 1;
+    for &v in adj[u].iter() {
+        if v != p {
+            calc_sz(adj, v, u, sub_sz);
+            sub_sz[u] += sub_sz[v];
+        }
+    }
+}
+
+fn dfs(
+    adj: &mut [Vec<usize>],
+    mut u: usize,
+    sub_sz: &mut [usize],
+    call_dfs: &mut dyn CentDecompDfs,
+) {
+    calc_sz(adj, u, u, sub_sz);
+    let sz_root = sub_sz[u];
+    let mut p = u;
+    loop {
+        let big_ch = adj[u]
+            .iter()
+            .filter(|&&v| v != p)
+            .find(|&&v| sub_sz[v] * 2 > sz_root);
+        if let Some(&v) = big_ch {
+            p = u;
+            u = v;
+        } else {
+            break;
+        }
+    }
+    call_dfs.dfs(adj, u);
+    for v in adj[u].clone() {
+        adj[v].retain(|&x| x != u);
+        dfs(adj, v, sub_sz, call_dfs);
+    }
+}
+
+/// A trait containing the DFS method which is called on each centroid of the tree with the
+/// back-edges outside of this centroid removed from `adj`
+pub trait CentDecompDfs {
+    /// The DFS method
+    fn dfs(&mut self, adj: &[Vec<usize>], cent: usize);
+}
+
+/// # Example
+/// - see count_paths_per_length.rs
+///
+/// # Params
+/// - `adj`: adjacency list representing an unrooted undirected tree
+/// - `call_dfs`: an object implementing `CentDecompDfs` trait
+///
+/// # Complexity
+/// - Time: O(n log n)
+/// - Space: O(n)
+pub fn cent_decomp(mut adj: Vec<Vec<usize>>, call_dfs: &mut dyn CentDecompDfs) {
+    let n = adj.len();
+    let mut sub_sz = vec![0; n];
+    for s in 0..n {
+        if sub_sz[s] == 0 {
+            dfs(&mut adj, s, &mut sub_sz, call_dfs);
+        }
+    }
+}

src/graphs/count_paths_per_length.rs

@@ -0,0 +1,96 @@
+//! # Count the number of paths of each length in a tree
+use crate::graphs::cent_decomp::{cent_decomp, CentDecompDfs};
+use crate::numbers::fft::fft_multiply;
+
+fn conv(a: &[u64], b: &[u64]) -> Vec<u64> {
+    let a = a.iter().map(|&x| x as f64).collect::<Vec<_>>();
+    let b = b.iter().map(|&x| x as f64).collect::<Vec<_>>();
+    let res_len = a.len() + b.len() - 1;
+    let mut res = fft_multiply(a, b);
+    res.resize(res_len, 0.0);
+    res.iter().map(|&x| x.round() as u64).collect()
+}
+
+struct CountPathsPerLength {
+    num_paths: Vec<u64>,
+}
+
+impl CountPathsPerLength {
+    fn new(n: usize) -> Self {
+        Self {
+            num_paths: vec![0; n],
+        }
+    }
+}
+
+impl CentDecompDfs for CountPathsPerLength {
+    fn dfs(&mut self, adj: &[Vec<usize>], cent: usize) {
+        let mut child_depths = vec![vec![]];
+        for &child in adj[cent].iter() {
+            let mut my_child_depths = vec![0];
+
+            use std::collections::VecDeque;
+
+            let mut q = VecDeque::new();
+            q.push_back((child, cent));
+
+            while !q.is_empty() {
+                my_child_depths.push(q.len() as u64);
+
+                let mut new_q = VecDeque::new();
+                while let Some((u, p)) = q.pop_front() {
+                    for &v in adj[u].iter() {
+                        if v != p {
+                            new_q.push_back((v, u));
+                        }
+                    }
+                }
+
+                q = new_q;
+            }
+
+            child_depths.push(my_child_depths);
+        }
+
+        child_depths.sort_by_key(|v| v.len());
+
+        let mut acc = vec![1];
+        for depth_arr in child_depths {
+            let res = conv(&acc, &depth_arr);
+            for (d, &cnt) in res.iter().enumerate() {
+                self.num_paths[d] += cnt;
+            }
+
+            acc.resize(acc.len().max(depth_arr.len()), 0);
+            for (d, &cnt) in depth_arr.iter().enumerate() {
+                acc[d] += cnt;
+            }
+        }
+    }
+}
+
+/// # Example
+/// ```
+/// use programming_team_code_rust::graphs::count_paths_per_length::count_paths_per_length;
+/// let adj = vec![
+///    vec![1, 2],
+///    vec![0, 3, 4],
+///    vec![0, 5, 6],
+///    vec![1],
+///    vec![1],
+///    vec![2],
+///    vec![2],
+/// ];
+/// let res = count_paths_per_length(&adj);
+/// assert_eq!(res, vec![0, 6, 7, 4, 4, 0, 0]);
+/// ```
+///
+/// # Complexity
+/// - Time: O(n log^2 n) with FFT for convolution
+/// - Space: O(n)
+pub fn count_paths_per_length(adj: &[Vec<usize>]) -> Vec<u64> {
+    let n = adj.len();
+    let mut obj = CountPathsPerLength::new(n);
+    cent_decomp(adj.to_vec(), &mut obj);
+    obj.num_paths
+}

src/graphs/mod.rs

@@ -1,4 +1,6 @@
 //! # Graph Algorithms
+pub mod cent_decomp;
+pub mod count_paths_per_length;
 mod dfs_order;
 pub mod dijk;
 pub mod lca;

src/numbers/fft.rs

@@ -0,0 +1,141 @@
+//! Fast Fourier Transform (FFT) implementation ripped from: <https://github.com/bminaiev/rust-contests/blob/main/algo_lib/src/math/fft.rs>
+use std::ops::{Add, Mul, MulAssign, Sub};
+
+#[derive(Copy, Clone)]
+struct Complex {
+    real: f64,
+    imag: f64,
+}
+
+impl Complex {
+    const ZERO: Self = Complex {
+        real: 0.0,
+        imag: 0.0,
+    };
+    const ONE: Self = Complex {
+        real: 1.0,
+        imag: 0.0,
+    };
+}
+
+impl Mul for Complex {
+    type Output = Complex;
+
+    fn mul(self, rhs: Self) -> Self::Output {
+        Self {
+            real: self.real * rhs.real - self.imag * rhs.imag,
+            imag: self.real * rhs.imag + self.imag * rhs.real,
+        }
+    }
+}
+
+impl MulAssign for Complex {
+    fn mul_assign(&mut self, rhs: Self) {
+        *self = *self * rhs;
+    }
+}
+
+impl Add for Complex {
+    type Output = Complex;
+
+    fn add(self, rhs: Self) -> Self::Output {
+        Self {
+            real: self.real + rhs.real,
+            imag: self.imag + rhs.imag,
+        }
+    }
+}
+
+impl Sub for Complex {
+    type Output = Complex;
+
+    fn sub(self, rhs: Self) -> Self::Output {
+        Self {
+            real: self.real - rhs.real,
+            imag: self.imag - rhs.imag,
+        }
+    }
+}
+
+fn fft(a: &mut [Complex], invert: bool) {
+    let n = a.len();
+    assert!(n.is_power_of_two());
+    let shift = usize::BITS - n.trailing_zeros();
+    for i in 1..n {
+        let j = (i << shift).reverse_bits();
+        assert!(j < n);
+        if i < j {
+            a.swap(i, j);
+        }
+    }
+    for len in (1..).map(|x| 1 << x).take_while(|s| *s <= n) {
+        let half = len / 2;
+        let alpha = std::f64::consts::PI * 2.0 / (len as f64);
+        let cos = f64::cos(alpha);
+        let sin = f64::sin(alpha) * (if invert { -1.0 } else { 1.0 });
+        let complex_angle = Complex {
+            real: cos,
+            imag: sin,
+        };
+        for start in (0..n).step_by(len) {
+            let mut mult = Complex::ONE;
+            for j in 0..half {
+                let u = a[start + j];
+                let v = a[start + half + j] * mult;
+                a[start + j] = u + v;
+                a[start + j + half] = u - v;
+                mult *= complex_angle;
+            }
+        }
+    }
+    if invert {
+        for elem in a.iter_mut().take(n) {
+            let n = n as f64;
+            elem.imag /= n;
+            elem.real /= n;
+        }
+    }
+}
+
+fn fft_multiply_raw(mut a: Vec<Complex>, mut b: Vec<Complex>) -> Vec<Complex> {
+    assert!(a.len().is_power_of_two());
+    assert!(b.len().is_power_of_two());
+    assert_eq!(a.len(), b.len());
+    fft(&mut a, false);
+    fft(&mut b, false);
+    for (x, y) in a.iter_mut().zip(b.iter()) {
+        *x *= *y;
+    }
+    fft(&mut a, true);
+    a
+}
+
+fn fft_multiply_complex(mut a: Vec<Complex>, mut b: Vec<Complex>) -> Vec<Complex> {
+    let expected_size = (a.len() + b.len() - 1).next_power_of_two();
+    a.resize(expected_size, Complex::ZERO);
+    b.resize(expected_size, Complex::ZERO);
+    fft_multiply_raw(a, b)
+}
+
+/// # Example
+/// ```
+/// use programming_team_code_rust::numbers::fft::fft_multiply;
+///
+/// let a = vec![1.0, 2.0, 3.0];
+/// let b = vec![4.0, 5.0, 6.0];
+/// let c = fft_multiply(a, b);
+/// let expected = vec![4.0, 13.0, 28.0, 27.0, 18.0];
+/// for (x, y) in c.iter().zip(expected.iter()) {
+///    assert!((x - y).abs() < 1e-15);
+/// }
+/// ```
+///
+/// # Complexity (n = max(a.len(), b.len()))
+/// Given for practical use cases, although consider it with a relatively large constant factor
+/// - Time: O(n log n)
+/// - Space: O(n)
+pub fn fft_multiply(a: Vec<f64>, b: Vec<f64>) -> Vec<f64> {
+    let a: Vec<_> = a.iter().map(|&x| Complex { real: x, imag: 0.0 }).collect();
+    let b: Vec<_> = b.iter().map(|&x| Complex { real: x, imag: 0.0 }).collect();
+    fft_multiply_complex(a, b).iter().map(|c| c.real).collect()
+}

src/numbers/mod.rs

@@ -1,3 +1,4 @@
 //! # Number Theory and Combinatorics
 pub mod binom;
+pub mod fft;
 pub mod primes;