centroid_decomposition/
lib.rs

1//! Centroid Decomposition of a tree.
2//!
3//! Centroid decomposition recursively splits a tree at its centroid —
4//! the vertex whose removal leaves components of size at most `n/2` —
5//! and repeats on each component. The resulting recursion tree has
6//! height O(log n), and any two original vertices share an ancestor
7//! in the decomposition at depth at most O(log n).
8//!
9//! This is the foundation for "small-to-large" tricks on trees, for
10//! answering distance queries offline (count paths of length ≤ k), and
11//! for many divide-and-conquer-on-tree problems.
12//!
13//! This crate exposes [`build_centroid_tree`], which returns the
14//! centroid tree: for each vertex, the centroid chosen at the next
15//! level up (its "centroid parent") and the depth of the vertex in the
16//! centroid tree.
17//!
18//! # Example
19//!
20//! ```
21//! use centroid_decomposition::build_centroid_tree;
22//!
23//! // Path 0 -- 1 -- 2 -- 3 -- 4
24//! let adj = vec![
25//!     vec![1],
26//!     vec![0, 2],
27//!     vec![1, 3],
28//!     vec![2, 4],
29//!     vec![3],
30//! ];
31//! let (parent, depth) = build_centroid_tree(&adj);
32//! // Centroid of a 5-path is the middle vertex (2).
33//! assert_eq!(parent[2], None);
34//! ```
35
36/// Build the centroid decomposition of the tree `adj`.
37///
38/// Returns two vectors of length `n`:
39/// * `parent[v]` — the centroid at the level above `v` in the
40///   decomposition (i.e., the centroid chosen for the component that
41///   contained `v` in the previous recursion step). `None` for the root
42///   of the centroid tree.
43/// * `depth[v]` — depth of `v` in the centroid tree (0 for the root).
44///
45/// # Panics
46/// Panics if `adj` is empty.
47pub fn build_centroid_tree(adj: &[Vec<usize>]) -> (Vec<Option<usize>>, Vec<usize>) {
48    let n = adj.len();
49    assert!(n > 0, "graph must be non-empty");
50
51    let mut parent: Vec<Option<usize>> = vec![None; n];
52    let mut depth = vec![0usize; n];
53    let mut removed = vec![false; n];
54    let mut size = vec![0usize; n];
55
56    decompose(adj, &mut removed, &mut size, 0, None, 0, &mut parent, &mut depth);
57
58    (parent, depth)
59}
60
61fn decompose(
62    adj: &[Vec<usize>],
63    removed: &mut [bool],
64    size: &mut [usize],
65    entry: usize,
66    centroid_parent: Option<usize>,
67    centroid_depth: usize,
68    parent: &mut [Option<usize>],
69    depth: &mut [usize],
70) {
71    // Compute subtree sizes within the current component, rooted at
72    // `entry`.
73    let root = entry;
74    compute_sizes(adj, removed, size, root, usize::MAX);
75    let total = size[root];
76
77    // Find centroid: a vertex `c` in the current component where every
78    // child component has size ≤ total/2.
79    let c = find_centroid(adj, removed, size, root, usize::MAX, total);
80
81    parent[c] = centroid_parent;
82    depth[c] = centroid_depth;
83    removed[c] = true;
84
85    // Recurse on each component formed by removing `c`.
86    for &u in &adj[c] {
87        if !removed[u] {
88            decompose(adj, removed, size, u, Some(c), centroid_depth + 1, parent, depth);
89        }
90    }
91    removed[c] = false;
92}
93
94fn compute_sizes(adj: &[Vec<usize>], removed: &[bool], size: &mut [usize], v: usize, p: usize) -> usize {
95    let mut s = 1usize;
96    for &u in &adj[v] {
97        if u != p && !removed[u] {
98            s += compute_sizes(adj, removed, size, u, v);
99        }
100    }
101    size[v] = s;
102    s
103}
104
105fn find_centroid(
106    adj: &[Vec<usize>],
107    removed: &[bool],
108    size: &[usize],
109    v: usize,
110    p: usize,
111    total: usize,
112) -> usize {
113    for &u in &adj[v] {
114        if u != p && !removed[u] && size[u] > total / 2 {
115            return find_centroid(adj, removed, size, u, v, total);
116        }
117    }
118    v
119}
centroid_decomposition/lib.rs

centroid_decomposition/
lib.rs
