competitive/tree/
euler_tour.rs

1use super::{Graph, RangeMinimumQuery, UndirectedSparseGraph};
2use std::{marker::PhantomData, ops::Range};
3
4pub trait EulerTourKind {
5    const USE_LAST: bool = false;
6    const USE_VISIT: bool = false;
7
8    fn size(n: usize) -> usize {
9        if Self::USE_VISIT {
10            2 * n - 1
11        } else if Self::USE_LAST {
12            2 * n
13        } else {
14            n
15        }
16    }
17}
18
19mod marker {
20    use super::EulerTourKind;
21
22    #[derive(Debug, Clone)]
23    pub enum First {}
24    #[derive(Debug, Clone)]
25    pub enum FirstLast {}
26    #[derive(Debug, Clone)]
27    pub enum Visit {}
28
29    impl EulerTourKind for First {}
30    impl EulerTourKind for FirstLast {
31        const USE_LAST: bool = true;
32    }
33    impl EulerTourKind for Visit {
34        const USE_VISIT: bool = true;
35    }
36}
37
38#[derive(Debug)]
39pub struct EulerTourBuilder<'a, K>
40where
41    K: EulerTourKind,
42{
43    tree: &'a UndirectedSparseGraph,
44    root: usize,
45    vidx: Vec<[usize; 2]>,
46    eidx: Vec<[usize; 2]>,
47    pos: usize,
48    _marker: PhantomData<fn() -> K>,
49}
50
51#[derive(Debug, Clone)]
52pub struct EulerTour<K>
53where
54    K: EulerTourKind,
55{
56    pub root: usize,
57    pub vidx: Vec<[usize; 2]>,
58    pub eidx: Vec<[usize; 2]>,
59    pub size: usize,
60    _marker: PhantomData<fn() -> K>,
61}
62
63impl<'a, K> EulerTourBuilder<'a, K>
64where
65    K: EulerTourKind,
66{
67    pub fn new(tree: &'a UndirectedSparseGraph, root: usize) -> Self {
68        let n = tree.vertices_size();
69        Self {
70            tree,
71            root,
72            vidx: vec![[0usize; 2]; n],
73            eidx: vec![[0usize; 2]; n - 1],
74            pos: 0,
75            _marker: PhantomData,
76        }
77    }
78
79    pub fn build_with_trace(mut self, mut trace: impl FnMut(usize)) -> EulerTour<K> {
80        self.dfs(self.root, !0, &mut trace);
81        EulerTour {
82            root: self.root,
83            vidx: self.vidx,
84            eidx: self.eidx,
85            size: self.pos,
86            _marker: PhantomData,
87        }
88    }
89
90    pub fn build(self) -> EulerTour<K> {
91        self.build_with_trace(|_u| {})
92    }
93
94    fn dfs(&mut self, u: usize, parent: usize, trace: &mut impl FnMut(usize)) {
95        self.vidx[u][0] = self.pos;
96        trace(u);
97        self.pos += 1;
98        for a in self.tree.neighbors(u) {
99            if a.to != parent {
100                self.eidx[a.label][0] = self.pos;
101                self.dfs(a.to, u, trace);
102                self.eidx[a.label][1] = self.pos;
103                if K::USE_VISIT {
104                    trace(u);
105                    self.pos += 1;
106                }
107            }
108        }
109        self.vidx[u][1] = self.pos;
110        if K::USE_LAST {
111            trace(u);
112            self.pos += 1;
113        }
114    }
115}
116
117impl EulerTourBuilder<'_, marker::First> {
118    pub fn build_with_rearrange<T>(self, s: &[T]) -> (EulerTour<marker::First>, Vec<T>)
119    where
120        T: Clone,
121    {
122        assert_eq!(s.len(), self.tree.vertices_size());
123        let mut trace = Vec::with_capacity(marker::First::size(s.len()));
124        let tour = self.build_with_trace(|u| {
125            trace.push(s[u].clone());
126        });
127        (tour, trace)
128    }
129}
130
131impl EulerTourBuilder<'_, marker::FirstLast> {
132    pub fn build_with_rearrange<T>(
133        self,
134        s: &[T],
135        mut inverse: impl FnMut(T) -> T,
136    ) -> (EulerTour<marker::FirstLast>, Vec<T>)
137    where
138        T: Clone,
139    {
140        assert_eq!(s.len(), self.tree.vertices_size());
141        let mut visited = vec![false; s.len()];
142        let mut trace = Vec::with_capacity(marker::FirstLast::size(s.len()));
143        let tour = self.build_with_trace(|u| {
144            if !visited[u] {
145                trace.push(s[u].clone());
146                visited[u] = true;
147            } else {
148                trace.push(inverse(s[u].clone()));
149            }
150        });
151        (tour, trace)
152    }
153}
154
155impl EulerTourBuilder<'_, marker::Visit> {
156    pub fn build_with_rearrange<T>(self, s: &[T]) -> (EulerTour<marker::Visit>, Vec<T>)
157    where
158        T: Clone,
159    {
160        assert_eq!(s.len(), self.tree.vertices_size());
161        let mut trace = Vec::with_capacity(marker::Visit::size(s.len()));
162        let tour = self.build_with_trace(|u| {
163            trace.push(s[u].clone());
164        });
165        (tour, trace)
166    }
167}
168
169impl UndirectedSparseGraph {
170    pub fn subtree_euler_tour_builder<'a>(
171        &'a self,
172        root: usize,
173    ) -> EulerTourBuilder<'a, marker::First> {
174        EulerTourBuilder::new(self, root)
175    }
176
177    pub fn path_euler_tour_builder<'a>(
178        &'a self,
179        root: usize,
180    ) -> EulerTourBuilder<'a, marker::FirstLast> {
181        EulerTourBuilder::new(self, root)
182    }
183
184    pub fn full_euler_tour_builder<'a>(
185        &'a self,
186        root: usize,
187    ) -> EulerTourBuilder<'a, marker::Visit> {
188        EulerTourBuilder::new(self, root)
189    }
190
191    pub fn lca(&self, root: usize) -> LowestCommonAncestor {
192        let (order, parents) = self.tree_order(root);
193        LowestCommonAncestor::from_dfs_preorder(&parents, &order)
194    }
195}
196
197impl EulerTour<marker::First> {
198    pub fn get<T>(&self, u: usize, mut f: impl FnMut(usize) -> T) -> T {
199        let [l, _] = self.vidx[u];
200        f(l)
201    }
202
203    pub fn update<T>(&self, u: usize, x: T, mut f: impl FnMut(usize, T)) {
204        let [l, _] = self.vidx[u];
205        f(l, x);
206    }
207
208    pub fn fold<T>(&self, u: usize, mut f: impl FnMut(Range<usize>) -> T) -> T {
209        let [l, r] = self.vidx[u];
210        f(l..r)
211    }
212
213    pub fn range_update<T>(&self, u: usize, x: T, mut f: impl FnMut(Range<usize>, T)) {
214        let [l, r] = self.vidx[u];
215        f(l..r, x);
216    }
217}
218
219impl EulerTour<marker::FirstLast> {
220    pub fn get<T>(&self, u: usize, mut f: impl FnMut(usize) -> T) -> T {
221        let [l, _] = self.vidx[u];
222        f(l)
223    }
224
225    pub fn update<T>(&self, u: usize, x: T, invx: T, mut f: impl FnMut(usize, T)) {
226        let [l, r] = self.vidx[u];
227        f(l, x);
228        f(r, invx);
229    }
230
231    // f: accumulate
232    pub fn fold<T>(&self, u: usize, mut f: impl FnMut(usize) -> T) -> T {
233        f(self.vidx[u][0])
234    }
235}
236
237#[derive(Debug)]
238pub struct LowestCommonAncestor {
239    node_to_index: Vec<u32>,
240    label_to_node: Vec<u32>,
241    rmq: RangeMinimumQuery<u32>,
242    depth: Vec<u32>,
243}
244
245impl LowestCommonAncestor {
246    /// `parents` must contain one parent per vertex and use `!0` for the root.
247    pub fn from_parents(parents: &[usize]) -> Self {
248        let n = parents.len();
249        let root = parents.iter().position(|&parent| parent == !0).unwrap();
250        let mut depth = vec![!0u32; n];
251        depth[root] = 0;
252        let mut label_to_node = Vec::with_capacity(n);
253        let mut node_to_label = vec![0u32; n];
254        label_to_node.push(root as u32);
255        let mut path = Vec::new();
256        for mut u in 0..n {
257            if depth[u] != !0 {
258                continue;
259            }
260            if depth[parents[u]] != !0 {
261                depth[u] = depth[parents[u]] + 1;
262                node_to_label[u] = label_to_node.len() as u32;
263                label_to_node.push(u as u32);
264                continue;
265            }
266            while depth[u] == !0 {
267                path.push(u);
268                u = parents[u];
269            }
270            while let Some(u) = path.pop() {
271                depth[u] = depth[parents[u]] + 1;
272                node_to_label[u] = label_to_node.len() as u32;
273                label_to_node.push(u as u32);
274            }
275        }
276
277        let mut label_to_index = vec![0u32; n];
278        for i in (1..n).rev() {
279            let u = label_to_node[i] as usize;
280            let parent = node_to_label[parents[u]] as usize;
281            label_to_index[parent] += label_to_index[i] + 1;
282        }
283        for i in 1..n {
284            let u = label_to_node[i] as usize;
285            let parent = node_to_label[parents[u]] as usize;
286            let descendants = label_to_index[i];
287            let next = label_to_index[parent];
288            label_to_index[i] = next;
289            label_to_index[parent] = next - descendants - 1;
290        }
291
292        let mut index_to_parent = vec![0u32; n];
293        for label in (1..n).rev() {
294            let u = label_to_node[label] as usize;
295            index_to_parent[label_to_index[label] as usize] = node_to_label[parents[u]];
296            node_to_label[u] = label_to_index[label];
297        }
298
299        Self {
300            node_to_index: node_to_label,
301            label_to_node,
302            rmq: RangeMinimumQuery::new(index_to_parent),
303            depth,
304        }
305    }
306
307    /// `order` must be a DFS preorder containing every vertex in `parents` exactly once.
308    /// `parents` uses `!0` for the root.
309    pub fn from_dfs_preorder(parents: &[usize], order: &[usize]) -> Self {
310        let n = parents.len();
311        let mut node_to_index = vec![0u32; n];
312        for (i, &u) in order.iter().enumerate() {
313            node_to_index[u] = i as u32;
314        }
315        let mut depth = vec![0u32; n];
316        let mut index_to_parent = vec![0u32; n];
317        for (i, &u) in order.iter().enumerate().skip(1) {
318            let p = parents[u];
319            depth[u] = depth[p] + 1;
320            index_to_parent[i] = node_to_index[p];
321        }
322        Self {
323            node_to_index,
324            label_to_node: order.iter().map(|&u| u as u32).collect(),
325            rmq: RangeMinimumQuery::new(index_to_parent),
326            depth,
327        }
328    }
329
330    #[inline]
331    pub fn depth(&self, u: usize) -> usize {
332        self.depth[u] as usize
333    }
334
335    #[inline]
336    pub fn lca(&self, u: usize, v: usize) -> usize {
337        if u == v {
338            return u;
339        }
340        let u = self.node_to_index[u] as usize;
341        let v = self.node_to_index[v] as usize;
342        let label = self.rmq.fold(u.min(v) + 1, u.max(v) + 1) as usize;
343        self.label_to_node[label] as usize
344    }
345}
346
347#[cfg(test)]
348mod tests {
349    use super::*;
350    use crate::{
351        algebra::{AdditiveOperation, RangeSumRangeAdd},
352        crecurse,
353        data_structure::{LazySegmentTree, SegmentTree},
354        tools::Xorshift,
355        tree::MixedTree,
356    };
357
358    #[test]
359    fn test_builder() {
360        let mut rng = Xorshift::default();
361        for _ in 0..200 {
362            let n = rng.random(1..=200);
363            let tree = rng.random(MixedTree(n));
364            let root = rng.random(0..n);
365            let et1 = tree.subtree_euler_tour_builder(root).build();
366            let et2 = tree.path_euler_tour_builder(root).build();
367            let et3 = tree.full_euler_tour_builder(root).build();
368            assert_eq!(et1.size, marker::First::size(n));
369            assert_eq!(et2.size, marker::FirstLast::size(n));
370            assert_eq!(et3.size, marker::Visit::size(n));
371            for u in 0..n {
372                assert!(et1.vidx[u][0] < et1.vidx[u][1]);
373                assert!(et1.vidx[u][1] <= marker::First::size(n));
374                assert!(et2.vidx[u][0] < et2.vidx[u][1]);
375                assert!(et2.vidx[u][1] < marker::FirstLast::size(n));
376                assert!(et3.vidx[u][0] < et3.vidx[u][1]);
377                assert!(et3.vidx[u][1] <= marker::Visit::size(n));
378            }
379        }
380    }
381
382    #[test]
383    fn test_subtree_euler_tour() {
384        const A: i64 = 1_000_000;
385        let mut rng = Xorshift::default();
386        for _ in 0..200 {
387            let n = rng.random(1..=200);
388            let tree = rng.random(MixedTree(n));
389            let root = rng.random(0..n);
390            let mut a: Vec<_> = rng.random_iter(0..A).take(n).collect();
391            let (et, arr) = tree
392                .subtree_euler_tour_builder(root)
393                .build_with_rearrange(&a);
394            let mut seg = LazySegmentTree::<RangeSumRangeAdd<i64>>::from_keys(arr.into_iter());
395            for _ in 0..200 {
396                match rng.random(0..4) {
397                    0 => {
398                        let u = rng.random(0..n);
399                        let result = et.get(u, |idx| seg.get(idx)).0;
400                        let expected = a[u];
401                        assert_eq!(result, expected);
402                    }
403                    1 => {
404                        let u = rng.random(0..n);
405                        let x = rng.random(0..A);
406                        et.update(u, x, |i, x| seg.update(i..=i, x));
407                        a[u] += x;
408                    }
409                    2 => {
410                        let u = rng.random(0..n);
411                        let result = et.fold(u, |r| seg.fold(r)).0;
412                        let mut expected = 0;
413                        crecurse!(
414                            unsafe fn dfs(v: usize, p: usize, b: bool) {
415                                let b = b || v == u;
416                                if b {
417                                    expected += a[v];
418                                }
419                                for a in tree.neighbors(v) {
420                                    if a.to != p {
421                                        dfs!(a.to, v, b);
422                                    }
423                                }
424                            }
425                        )(root, !0, false);
426                        assert_eq!(result, expected);
427                    }
428                    _ => {
429                        let u = rng.random(0..n);
430                        let x = rng.random(0..A);
431                        et.range_update(u, x, |r, x| seg.update(r, x));
432                        crecurse!(
433                            unsafe fn dfs(v: usize, p: usize, b: bool) {
434                                let b = b || v == u;
435                                if b {
436                                    a[v] += x;
437                                }
438                                for a in tree.neighbors(v) {
439                                    if a.to != p {
440                                        dfs!(a.to, v, b);
441                                    }
442                                }
443                            }
444                        )(root, !0, false);
445                    }
446                }
447            }
448        }
449    }
450
451    #[test]
452    fn test_path_euler_tour() {
453        const A: i64 = 1_000_000;
454        let mut rng = Xorshift::default();
455        for _ in 0..200 {
456            let n = rng.random(1..=200);
457            let tree = rng.random(MixedTree(n));
458            let root = rng.random(0..n);
459            let mut a: Vec<_> = rng.random_iter(0..A).take(n).collect();
460            let (et, arr) = tree
461                .path_euler_tour_builder(root)
462                .build_with_rearrange(&a, |x| -x);
463            let mut seg = SegmentTree::<AdditiveOperation<i64>>::from_vec(arr);
464            for _ in 0..200 {
465                match rng.random(0..3) {
466                    0 => {
467                        let u = rng.random(0..n);
468                        let result = et.get(u, |idx| seg.get(idx));
469                        let expected = a[u];
470                        assert_eq!(result, expected);
471                    }
472                    1 => {
473                        let u = rng.random(0..n);
474                        let x = rng.random(0..A);
475                        let invx = -x;
476                        et.update(u, x, invx, |i, x| seg.update(i, x));
477                        a[u] += x;
478                    }
479                    _ => {
480                        let u = rng.random(0..n);
481                        let result = et.fold(u, |k| seg.fold(0..=k));
482                        let mut expected = 0;
483                        crecurse!(
484                            unsafe fn dfs(v: usize, p: usize) -> bool {
485                                if v == u {
486                                    expected += a[v];
487                                    return true;
488                                }
489                                for adj in tree.neighbors(v) {
490                                    if adj.to != p && dfs!(adj.to, v) {
491                                        expected += a[v];
492                                        return true;
493                                    }
494                                }
495                                false
496                            }
497                        )(root, !0);
498                        assert_eq!(result, expected);
499                    }
500                }
501            }
502        }
503    }
504
505    #[test]
506    fn test_lca() {
507        let mut rng = Xorshift::default();
508        for _ in 0..200 {
509            let n = rng.random(1..=200);
510            let tree = rng.random(MixedTree(n));
511            let root = rng.random(0..n);
512            let (_, parents) = tree.tree_order(root);
513            let lcas = [tree.lca(root), LowestCommonAncestor::from_parents(&parents)];
514            for _ in 0..200 {
515                let u = rng.random(0..n);
516                let v = rng.random(0..n);
517                let expected = crecurse!(
518                    unsafe fn dfs(w: usize, p: usize) -> Result<usize, [bool; 2]> {
519                        let mut found = [false; 2];
520                        if w == u {
521                            found[0] = true;
522                        }
523                        if w == v {
524                            found[1] = true;
525                        }
526                        for adj in tree.neighbors(w) {
527                            if adj.to != p {
528                                match dfs!(adj.to, w) {
529                                    Ok(lca) => return Ok(lca),
530                                    Err(res) => {
531                                        for i in 0..2 {
532                                            if res[i] {
533                                                found[i] = true;
534                                            }
535                                        }
536                                    }
537                                }
538                            }
539                        }
540                        if found[0] && found[1] {
541                            Ok(w)
542                        } else {
543                            Err(found)
544                        }
545                    }
546                )(root, !0)
547                .unwrap();
548                for lca in &lcas {
549                    assert_eq!(lca.lca(u, v), expected);
550                }
551            }
552        }
553    }
554}
competitive/tree/euler_tour.rs

competitive/tree/
euler_tour.rs
