competitive/tree/
level_ancestor.rs

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

competitive/tree/
level_ancestor.rs
