competitive/tree/
level_ancestor.rs

1use super::UndirectedSparseGraph;
2
3pub struct LevelAncestor {
4    parent: Vec<usize>,
5    depth: Vec<usize>,
6    start: Vec<usize>,
7    index: Vec<usize>,
8    ladder: Vec<usize>,
9}
10
11impl UndirectedSparseGraph {
12    pub fn level_ancestor(&self, root: usize) -> LevelAncestor {
13        let n = self.vertices_size();
14        let (order, parent) = self.tree_order(root);
15        let mut depth = vec![0; n];
16        for &u in order.iter().skip(1) {
17            depth[u] = depth[parent[u]] + 1;
18        }
19        let mut height = vec![1; n];
20        let mut heavy = vec![n; n];
21        for &u in order.iter().skip(1).rev() {
22            let p = parent[u];
23            if heavy[p] == n || height[heavy[p]] < height[u] {
24                heavy[p] = u;
25            }
26            height[p] = height[p].max(height[u] + 1);
27        }
28
29        let mut start = vec![0; n];
30        let mut index = vec![0; n];
31        let mut ladder = Vec::with_capacity(2 * n);
32        for &head in &order {
33            if head != root && heavy[parent[head]] == head {
34                continue;
35            }
36            let extension = height[head].min(depth[head]);
37            let offset = ladder.len();
38            ladder.resize(offset + extension + height[head], n);
39            let mut u = head;
40            for i in (0..extension).rev() {
41                u = parent[u];
42                ladder[offset + i] = u;
43            }
44            let mut u = head;
45            for i in extension..extension + height[head] {
46                ladder[offset + i] = u;
47                start[u] = offset;
48                index[u] = offset + i;
49                u = heavy[u];
50            }
51        }
52
53        LevelAncestor {
54            parent,
55            depth,
56            start,
57            index,
58            ladder,
59        }
60    }
61
62    pub fn level_ancestor_batch(
63        &self,
64        root: usize,
65        queries: impl IntoIterator<Item = (usize, usize)>,
66    ) -> Vec<Option<usize>> {
67        let n = self.vertices_size();
68        let mut start = vec![0; n + 1];
69        let queries: Vec<(usize, usize)> = queries.into_iter().collect();
70        for &(u, _) in &queries {
71            start[u] += 1;
72        }
73        for d in 0..n {
74            start[d + 1] += start[d];
75        }
76        let qsize = queries.len();
77        let mut batch = vec![(0, 0); qsize];
78        for (i, &(u, k)) in queries.iter().enumerate() {
79            start[u] -= 1;
80            batch[start[u]] = (k, i);
81        }
82        let (order, parent) = self.tree_order(root);
83        let mut path = Vec::with_capacity(n);
84        let mut results = vec![None; qsize];
85        for u in order {
86            while path.last().is_some_and(|&v| v != parent[u]) {
87                path.pop();
88            }
89            path.push(u);
90            for &(k, qi) in &batch[start[u]..start[u + 1]] {
91                let depth = path.len() - 1;
92                if k <= depth {
93                    results[qi] = Some(path[depth - k]);
94                }
95            }
96        }
97        results
98    }
99}
100
101impl LevelAncestor {
102    #[inline]
103    pub fn la(&self, mut u: usize, mut k: usize) -> Option<usize> {
104        if self.depth[u] < k {
105            return None;
106        }
107        loop {
108            let start = self.start[u];
109            let index = self.index[u];
110            if k <= index - start {
111                return Some(self.ladder[index - k]);
112            }
113            k -= index - start + 1;
114            u = self.parent[self.ladder[start]];
115        }
116    }
117
118    #[inline]
119    pub fn depth(&self, u: usize) -> usize {
120        self.depth[u]
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use crate::{graph::Graph, tools::Xorshift, tree::MixedTree};
127
128    #[test]
129    fn test_level_ancestor() {
130        let mut rng = Xorshift::default();
131        for _ in 0..500 {
132            let n = rng.random(1..=200);
133            let tree = rng.random(MixedTree(n));
134            let root = rng.random(0..n);
135            let la = tree.level_ancestor(root);
136            let mut parent = vec![None; n];
137            let mut stack = vec![(root, None)];
138            while let Some((u, p)) = stack.pop() {
139                parent[u] = p;
140                for a in tree.neighbors(u) {
141                    if Some(a.to) != p {
142                        stack.push((a.to, Some(u)));
143                    }
144                }
145            }
146            let mut queries = vec![];
147            let mut results = vec![];
148            for u in 0..n {
149                let mut v = Some(u);
150                for d in 0..=n {
151                    assert_eq!(la.la(u, d), v);
152                    queries.push((u, d));
153                    results.push(v);
154                    v = v.and_then(|x| parent[x]);
155                }
156            }
157            assert_eq!(tree.level_ancestor_batch(root, queries), results);
158        }
159    }
160}
competitive/tree/level_ancestor.rs

competitive/tree/
level_ancestor.rs
