competitive/algorithm/
cartesian_tree.rs1use 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}