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 stack = vec![];
18        for i in 0..a.len() {
19            let mut prev = !0usize;
20            while let Some(last) = stack.pop_if(|last| a[i] < a[*last]) {
21                prev = last;
22            }
23            if prev != !0 {
24                parents[prev] = i;
25            }
26            if let Some(&last) = stack.last() {
27                parents[i] = last;
28            }
29            stack.push(i);
30        }
31        let mut root = !0;
32        for i in 0..a.len() {
33            if parents[i] != !0 {
34                children[parents[i]][(i > parents[i]) as usize] = i;
35            } else {
36                root = i;
37            }
38        }
39        Self {
40            root,
41            parents,
42            children,
43        }
44    }
45    pub fn with_ranges(&self, mut f: impl FnMut(usize, Range<usize>)) {
46        let mut stack = vec![(self.root, 0, self.parents.len())];
47        while let Some((v, l, r)) = stack.pop() {
48            f(v, l..r);
49            if self.children[v][1] != !0 {
50                stack.push((self.children[v][1], v + 1, r));
51            }
52            if self.children[v][0] != !0 {
53                stack.push((self.children[v][0], l, v));
54            }
55        }
56    }
57}
58
59#[cfg(test)]
60mod tests {
61    use super::*;
62    use crate::{
63        algebra::MinOperation, crecurse, data_structure::SegmentTree, rand, tools::Xorshift,
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        for _ in 0..Q {
73            rand!(rng, n: 1..=N, a: [0..A; n]);
74            let mut seg = SegmentTree::<MinOperation<_>>::from_vec(
75                a.iter().enumerate().map(|(i, &a)| (a, i)).collect(),
76            );
77            let mut parents = vec![!0; n];
78            let mut root = !0;
79            let mut children = vec![[!0; 2]; n];
80            let mut ranges = vec![];
81            crecurse!(
82                unsafe fn dfs(l: usize, r: usize, p: usize, ci: usize) {
83                    if l >= r {
84                        return;
85                    }
86                    let m = seg.fold(l..r).1;
87                    ranges.push((m, l..r));
88                    if p == !0 {
89                        root = m;
90                    } else {
91                        parents[m] = p;
92                        children[p][ci] = m;
93                    }
94                    seg.set(m, (i64::MAX, usize::MAX));
95                    dfs!(l, m, m, 0);
96                    dfs!(m + 1, r, m, 1);
97                }
98            )(0, n, !0, !0);
99
100            let ct = CartesianTree::new(&a);
101            assert_eq!(ct.root, root);
102            assert_eq!(ct.parents, parents);
103            assert_eq!(ct.children, children);
104            let mut ct_ranges = vec![];
105            ct.with_ranges(|v, range| {
106                ct_ranges.push((v, range));
107            });
108            assert_eq!(ct_ranges, ranges);
109        }
110    }
111}