heavy_light_decomposition/
lib.rs

1//! Heavy-Light Decomposition (HLD).
2//!
3//! Heavy-Light Decomposition splits a tree into "heavy" paths (chains of
4//! vertices where each parent has its largest subtree as its first
5//! child) and "light" edges (connecting chains). Once decomposed, any
6//! path from `u` to `v` can be split into O(log n) segments, each lying
7//! within a single chain. By storing per-chain data in a Fenwick tree
8//! or segment tree, you can answer path queries / updates in
9//! O(log² n) overall.
10//!
11//! This crate ships the decomposition itself: it computes for each
12//! vertex its parent, depth, subtree size, the head of the chain it
13//! belongs to, and its position in the linearised chain order. You can
14//! then plug in any range-query data structure of your choice on top
15//! of that linearisation.
16//!
17//! # Example
18//!
19//! ```
20//! use heavy_light_decomposition::Hld;
21//!
22//! // Tree:
23//! //      0
24//! //      |
25//! //      1
26//! //     / \
27//! //    2   3
28//! //    |   |
29//! //    4   5
30//! let adj = vec![
31//!     vec![1],
32//!     vec![0, 2, 3],
33//!     vec![1, 4],
34//!     vec![1, 5],
35//!     vec![2],
36//!     vec![3],
37//! ];
38//! let hld = Hld::new(&adj, 0);
39//! assert_eq!(hld.parent[1], Some(0));
40//! assert_eq!(hld.depth[4], 3);
41//! ```
42
43/// Heavy-Light Decomposition over a rooted tree.
44#[derive(Debug, Clone)]
45pub struct Hld {
46    /// Parent of each vertex in the rooted tree. `None` for the root.
47    pub parent: Vec<Option<usize>>,
48    /// Depth of each vertex (root has depth 0).
49    pub depth: Vec<usize>,
50    /// Subtree size of each vertex (including itself).
51    pub size: Vec<usize>,
52    /// Head of the heavy chain containing each vertex.
53    pub head: Vec<usize>,
54    /// Position of each vertex in the linearised order of chains.
55    pub pos: Vec<usize>,
56}
57
58impl Hld {
59    /// Build the decomposition for the tree `adj` rooted at `root`.
60    ///
61    /// `adj[v]` is the list of neighbours of `v`. The tree is assumed
62    /// to be connected and undirected; edges to a vertex's parent are
63    /// filtered automatically.
64    ///
65    /// # Panics
66    /// Panics if `root >= adj.len()`.
67    pub fn new(adj: &[Vec<usize>], root: usize) -> Self {
68        let n = adj.len();
69        assert!(root < n, "root vertex out of range");
70
71        // Phase 1: parent + depth + subtree size via iterative DFS.
72        let mut parent: Vec<Option<usize>> = vec![None; n];
73        let mut depth = vec![0usize; n];
74        let mut size = vec![1usize; n];
75        // We store (vertex, parent, state) where state = 0 means "enter"
76        // and state = 1 means "exit".
77        let mut stack: Vec<(usize, Option<usize>, u8)> = vec![(root, None, 0)];
78        let mut order: Vec<usize> = Vec::with_capacity(n);
79        while let Some((v, p, state)) = stack.pop() {
80            if state == 0 {
81                parent[v] = p;
82                depth[v] = match p {
83                    Some(pv) => depth[pv] + 1,
84                    None => 0,
85                };
86                stack.push((v, p, 1));
87                for &u in &adj[v] {
88                    if Some(u) != p {
89                        stack.push((u, Some(v), 0));
90                    }
91                }
92            } else {
93                order.push(v);
94            }
95        }
96        for &v in &order {
97            if v != root {
98                if let Some(p) = parent[v] {
99                    size[p] += size[v];
100                }
101            }
102        }
103
104        // Phase 2: chain heads + positions via DFS from root, picking
105        // the heavy child as the next link.
106        let mut head = vec![root; n];
107        let mut pos = vec![0usize; n];
108        let mut cur_pos = 0usize;
109        dfs_decompose(adj, root, usize::MAX, root, &size, &mut head, &mut pos, &mut cur_pos);
110
111        Hld {
112            parent,
113            depth,
114            size,
115            head,
116            pos,
117        }
118    }
119
120    /// Decompose the path between `u` and `v` into half-open segments
121    /// `[lo, hi)` in the linearised chain order. Pass these segments to
122    /// a Fenwick / segment tree built on `pos` to answer path queries.
123    pub fn path_segments(&self, u: usize, v: usize) -> Vec<(usize, usize)> {
124        let mut segs = Vec::new();
125        let (mut u, mut v) = (u, v);
126        while self.head[u] != self.head[v] {
127            if self.depth[self.head[u]] > self.depth[self.head[v]] {
128                let h = self.head[u];
129                segs.push((self.pos[h], self.pos[u] + 1));
130                u = self.parent[h].unwrap();
131            } else {
132                let h = self.head[v];
133                segs.push((self.pos[h], self.pos[v] + 1));
134                v = self.parent[h].unwrap();
135            }
136        }
137        let (lo, hi) = if self.pos[u] <= self.pos[v] {
138            (self.pos[u], self.pos[v] + 1)
139        } else {
140            (self.pos[v], self.pos[u] + 1)
141        };
142        segs.push((lo, hi));
143        segs
144    }
145}
146
147fn dfs_decompose(
148    adj: &[Vec<usize>],
149    v: usize,
150    p: usize,
151    head: usize,
152    size: &[usize],
153    head_out: &mut [usize],
154    pos_out: &mut [usize],
155    cur_pos: &mut usize,
156) {
157    head_out[v] = head;
158    pos_out[v] = *cur_pos;
159    *cur_pos += 1;
160
161    // Find the heavy child (largest subtree).
162    let mut heavy: Option<usize> = None;
163    let mut heavy_size = 0usize;
164    for &u in &adj[v] {
165        if u != p && size[u] > heavy_size {
166            heavy_size = size[u];
167            heavy = Some(u);
168        }
169    }
170
171    // Recurse into the heavy child continuing the chain.
172    if let Some(h) = heavy {
173        dfs_decompose(adj, h, v, head, size, head_out, pos_out, cur_pos);
174    }
175
176    // Recurse into light children starting new chains.
177    for &u in &adj[v] {
178        if u != p && Some(u) != heavy {
179            dfs_decompose(adj, u, v, u, size, head_out, pos_out, cur_pos);
180        }
181    }
182}
heavy_light_decomposition/lib.rs

heavy_light_decomposition/
lib.rs
