Skip to main content

competitive/algorithm/
cartesian_tree.rs

1use std::ops::Range;
2
3#[derive(Debug, Clone)]
4pub struct CartesianTree {
5    pub root: usize,
6    pub parents: Vec<usize>,
7    pub children: Vec<[usize; 2]>,
8}
9
10impl CartesianTree {
11    pub fn new<T>(a: &[T]) -> Self
12    where
13        T: PartialOrd,
14    {
15        let mut parents = vec![!0; a.len()];
16        let mut children = vec![[!0; 2]; a.len()];
17        let mut root = !0;
18        for i in 0..a.len() {
19            let mut prev = !0;
20            let mut parent = i.wrapping_sub(1);
21            while parent != !0 && a[i] < a[parent] {
22                prev = parent;
23                parent = parents[parent];
24            }
25            parents[i] = parent;
26            if parent != !0 {
27                children[parent][1] = i;
28            } else {
29                root = i;
30            }
31            if prev != !0 {
32                parents[prev] = i;
33                children[i][0] = prev;
34            }
35        }
36        Self {
37            root,
38            parents,
39            children,
40        }
41    }
42    pub fn with_ranges(&self, mut f: impl FnMut(usize, Range<usize>)) {
43        let mut stack = vec![(self.root, 0, self.parents.len())];
44        while let Some((v, l, r)) = stack.pop() {
45            f(v, l..r);
46            if self.children[v][1] != !0 {
47                stack.push((self.children[v][1], v + 1, r));
48            }
49            if self.children[v][0] != !0 {
50                stack.push((self.children[v][0], l, v));
51            }
52        }
53    }
54}
55
56#[cfg(test)]
57mod tests {
58    use super::*;
59    use crate::{
60        algebra::MinOperation,
61        crecurse,
62        data_structure::SegmentTree,
63        tools::{Xorshift, testutil::exhaustive_sequences},
64    };
65
66    #[test]
67    fn test_cartesian_tree() {
68        const Q: usize = 1000;
69        const N: usize = 100;
70        const A: i64 = 100;
71        let mut rng = Xorshift::default();
72        let arrays = exhaustive_sequences(-1..=1, 0..=6).chain((0..Q).map(|_| {
73            let n = rng.random(1..=N);
74            rng.random_iter(0..A).take(n).collect()
75        }));
76        for a in arrays {
77            let n = a.len();
78            let mut seg = SegmentTree::<MinOperation<_>>::from_vec(
79                a.iter().enumerate().map(|(i, &a)| (a, i)).collect(),
80            );
81            let mut parents = vec![!0; n];
82            let mut root = !0;
83            let mut children = vec![[!0; 2]; n];
84            let mut ranges = vec![];
85            crecurse!(
86                unsafe fn dfs(l: usize, r: usize, p: usize, ci: usize) {
87                    if l >= r {
88                        return;
89                    }
90                    let m = seg.fold(l..r).1;
91                    ranges.push((m, l..r));
92                    if p == !0 {
93                        root = m;
94                    } else {
95                        parents[m] = p;
96                        children[p][ci] = m;
97                    }
98                    seg.set(m, (i64::MAX, usize::MAX));
99                    dfs!(l, m, m, 0);
100                    dfs!(m + 1, r, m, 1);
101                }
102            )(0, n, !0, !0);
103
104            let ct = CartesianTree::new(&a);
105            assert_eq!(ct.root, root);
106            assert_eq!(ct.parents, parents);
107            assert_eq!(ct.children, children);
108            let mut ct_ranges = vec![];
109            if !a.is_empty() {
110                ct.with_ranges(|v, range| {
111                    ct_ranges.push((v, range));
112                });
113            }
114            assert_eq!(ct_ranges, ranges);
115        }
116    }
117}