cp_library_rs/data_structure/
cartesian_tree.rs

1use std::marker::PhantomData;
2
3/// CartesianTree
4///
5/// inorder が元配列の順になる(二分木としての中順走査が 0..n-1)
6///
7/// `build_max` は max-heap(親の値 \(\ge\) 子の値).
8/// `build_min` は min-heap(親の値 \(\le\) 子の値).
9#[derive(Debug, Clone)]
10pub struct CartesianTree<T> {
11    pub n: usize,
12    pub root: usize,
13    pub parent: Vec<Option<usize>>,
14    pub left: Vec<Option<usize>>,
15    pub right: Vec<Option<usize>>,
16    _phantom: PhantomData<T>,
17}
18
19impl<T: Ord> CartesianTree<T> {
20    /// 共通実装.
21    ///
22    /// `should_pop(top, cur)` が true の間,スタックを pop する.
23    pub fn build_by<F>(arr: &[T], mut should_pop: F) -> Self
24    where
25        F: FnMut(&T, &T) -> bool,
26    {
27        let n = arr.len();
28        assert!(n > 0);
29
30        let mut parent = vec![None; n];
31        let mut left = vec![None; n];
32        let mut right = vec![None; n];
33
34        let mut st: Vec<usize> = Vec::with_capacity(n);
35
36        for i in 0..n {
37            let mut last: Option<usize> = None;
38
39            while let Some(&j) = st.last() {
40                if !should_pop(&arr[j], &arr[i]) {
41                    break;
42                }
43                last = Some(st.pop().unwrap());
44            }
45
46            // 最後に pop された `last` が `i` の左部分木(左の子の根)
47            if let Some(l) = last {
48                parent[l] = Some(i);
49                left[i] = Some(l);
50            }
51
52            // 残ったスタック top の右の子が `i`
53            if let Some(&top) = st.last() {
54                parent[i] = Some(top);
55                right[top] = Some(i);
56            }
57
58            st.push(i);
59        }
60
61        let root = (0..n).find(|&i| parent[i].is_none()).unwrap();
62
63        Self {
64            n,
65            root,
66            parent,
67            left,
68            right,
69            _phantom: PhantomData,
70        }
71    }
72
73    /// max-heap の Cartesian Tree を構築する(同値は右側優先).
74    ///
75    /// すなわち,区間の最大値(同値なら最右)を根にして再帰分割した木になる.
76    pub fn build_max(arr: &[T]) -> Self {
77        // stack top `top` が `cur` より「優先度が低い」なら pop する
78        // max-heap かつ同値は右優先なので,`top <= cur` の間 pop
79        Self::build_by(arr, |top, cur| top <= cur)
80    }
81
82    /// min-heap の Cartesian Tree を構築する(同値は右側優先).
83    ///
84    /// すなわち,区間の最小値(同値なら最右)を根にして再帰分割した木になる.
85    pub fn build_min(arr: &[T]) -> Self {
86        // min-heap かつ同値は右優先なので,`top >= cur` の間 pop
87        Self::build_by(arr, |top, cur| top >= cur)
88    }
89}