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.last() {
21                if a[i] < a[last] {
22                    prev = stack.pop().unwrap();
23                } else {
24                    break;
25                }
26            }
27            if prev != !0 {
28                parents[prev] = i;
29            }
30            if let Some(&last) = stack.last() {
31                parents[i] = last;
32            }
33            stack.push(i);
34        }
35        let mut root = !0;
36        for i in 0..a.len() {
37            if parents[i] != !0 {
38                children[parents[i]][(i > parents[i]) as usize] = i;
39            } else {
40                root = i;
41            }
42        }
43        Self {
44            root,
45            parents,
46            children,
47        }
48    }
49    pub fn with_ranges(&self, mut f: impl FnMut(usize, Range<usize>)) {
50        let mut stack = vec![(self.root, 0, self.parents.len())];
51        while let Some((v, l, r)) = stack.pop() {
52            f(v, l..r);
53            if self.children[v][1] != !0 {
54                stack.push((self.children[v][1], v + 1, r));
55            }
56            if self.children[v][0] != !0 {
57                stack.push((self.children[v][0], l, v));
58            }
59        }
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66    use crate::{
67        algebra::MinOperation, crecurse, data_structure::SegmentTree, rand, tools::Xorshift,
68    };
69
70    #[test]
71    fn test_cartesian_tree() {
72        const Q: usize = 1000;
73        const N: usize = 100;
74        const A: i64 = 100;
75        let mut rng = Xorshift::default();
76        for _ in 0..Q {
77            rand!(rng, n: 1..=N, a: [0..A; n]);
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            ct.with_ranges(|v, range| {
110                ct_ranges.push((v, range));
111            });
112            assert_eq!(ct_ranges, ranges);
113        }
114    }
115}
competitive/algorithm/cartesian_tree.rs

competitive/algorithm/
cartesian_tree.rs
