competitive/data_structure/binary_search_tree/
seeker.rs

1use super::{BstDataAccess, BstImmutRef, BstSpec, LazyMapMonoid, data, data::LazyMapElement};
2use std::{borrow::Borrow, cmp::Ordering, marker::PhantomData};
3
4pub trait BstSeeker {
5    type Spec: BstSpec;
6
7    fn bst_seek(&mut self, _node: BstImmutRef<'_, Self::Spec>) -> Ordering;
8}
9
10pub struct SeekLeft<Spec> {
11    _marker: PhantomData<fn() -> Spec>,
12}
13
14impl<S> Default for SeekLeft<S> {
15    fn default() -> Self {
16        Self {
17            _marker: PhantomData,
18        }
19    }
20}
21
22impl<Spec> BstSeeker for SeekLeft<Spec>
23where
24    Spec: BstSpec,
25{
26    type Spec = Spec;
27
28    fn bst_seek(&mut self, node: BstImmutRef<'_, Self::Spec>) -> Ordering {
29        if node.reborrow().right().descend().is_ok() {
30            Ordering::Greater
31        } else {
32            Ordering::Equal
33        }
34    }
35}
36
37pub struct SeekRight<Spec> {
38    _marker: PhantomData<fn() -> Spec>,
39}
40
41impl<S> Default for SeekRight<S> {
42    fn default() -> Self {
43        Self {
44            _marker: PhantomData,
45        }
46    }
47}
48
49impl<Spec> BstSeeker for SeekRight<Spec>
50where
51    Spec: BstSpec,
52{
53    type Spec = Spec;
54    fn bst_seek(&mut self, node: BstImmutRef<'_, Self::Spec>) -> Ordering {
55        if node.reborrow().left().descend().is_ok() {
56            Ordering::Less
57        } else {
58            Ordering::Equal
59        }
60    }
61}
62
63pub struct SeekByKey<'a, Spec, K, Q>
64where
65    Q: ?Sized,
66{
67    key: &'a Q,
68    _marker: PhantomData<fn() -> (Spec, K)>,
69}
70
71impl<'a, Spec, K, Q> SeekByKey<'a, Spec, K, Q>
72where
73    Q: ?Sized,
74{
75    pub fn new(key: &'a Q) -> Self {
76        Self {
77            key,
78            _marker: PhantomData,
79        }
80    }
81}
82
83impl<Spec, K, Q> BstSeeker for SeekByKey<'_, Spec, K, Q>
84where
85    Spec: BstSpec<Data: BstDataAccess<data::marker::Key, Value = K>>,
86    K: Borrow<Q>,
87    Q: Ord + ?Sized,
88{
89    type Spec = Spec;
90
91    fn bst_seek(&mut self, node: BstImmutRef<'_, Self::Spec>) -> Ordering {
92        node.reborrow()
93            .into_data()
94            .bst_data()
95            .borrow()
96            .cmp(self.key)
97    }
98}
99
100pub struct SeekBySize<Spec> {
101    index: usize,
102    _marker: PhantomData<fn() -> Spec>,
103}
104
105impl<Spec> SeekBySize<Spec> {
106    pub fn new(index: usize) -> Self {
107        Self {
108            index,
109            _marker: PhantomData,
110        }
111    }
112}
113
114impl<Spec> BstSeeker for SeekBySize<Spec>
115where
116    Spec: BstSpec<Data: BstDataAccess<data::marker::Size, Value = usize>>,
117{
118    type Spec = Spec;
119
120    fn bst_seek(&mut self, node: BstImmutRef<'_, Self::Spec>) -> Ordering {
121        let lsize = node
122            .reborrow()
123            .left()
124            .descend()
125            .map(|l| *l.into_data().bst_data())
126            .unwrap_or_default();
127        let ord = lsize.cmp(&self.index);
128        if matches!(ord, Ordering::Less) {
129            self.index -= lsize + 1;
130        }
131        ord
132    }
133}
134
135pub struct SeekByAccCond<Spec, L, F>
136where
137    L: LazyMapMonoid,
138{
139    acc: L::Agg,
140    f: F,
141    _marker: PhantomData<fn() -> (Spec, L)>,
142}
143
144impl<Spec, L, F> SeekByAccCond<Spec, L, F>
145where
146    L: LazyMapMonoid,
147    F: FnMut(&L::Agg) -> bool,
148{
149    pub fn new(f: F) -> Self {
150        Self {
151            acc: L::agg_unit(),
152            f,
153            _marker: PhantomData,
154        }
155    }
156}
157
158impl<Spec, L, F> BstSeeker for SeekByAccCond<Spec, L, F>
159where
160    Spec: BstSpec<Data: BstDataAccess<data::marker::LazyMap, Value = LazyMapElement<L>>>,
161    L: LazyMapMonoid,
162    F: FnMut(&L::Agg) -> bool,
163{
164    type Spec = Spec;
165
166    fn bst_seek(&mut self, node: BstImmutRef<'_, Self::Spec>) -> Ordering {
167        if let Ok(left) = node.reborrow().left().descend() {
168            let left_agg = &left.into_data().bst_data().agg;
169            let nagg = L::agg_operate(&self.acc, left_agg);
170            if (self.f)(&nagg) {
171                return Ordering::Greater;
172            }
173            let nagg = L::agg_operate(
174                &nagg,
175                &L::single_agg(&node.reborrow().into_data().bst_data().key),
176            );
177            if (self.f)(&nagg) {
178                Ordering::Equal
179            } else {
180                self.acc = nagg;
181                Ordering::Less
182            }
183        } else {
184            let nagg = L::agg_operate(
185                &self.acc,
186                &L::single_agg(&node.reborrow().into_data().bst_data().key),
187            );
188            if (self.f)(&nagg) {
189                Ordering::Equal
190            } else {
191                self.acc = nagg;
192                Ordering::Less
193            }
194        }
195    }
196}
197
198pub struct SeekByRaccCond<Spec, L, F>
199where
200    L: LazyMapMonoid,
201{
202    acc: L::Agg,
203    f: F,
204    _marker: PhantomData<fn() -> (Spec, L)>,
205}
206
207impl<Spec, L, F> SeekByRaccCond<Spec, L, F>
208where
209    L: LazyMapMonoid,
210    F: FnMut(&L::Agg) -> bool,
211{
212    pub fn new(f: F) -> Self {
213        Self {
214            acc: L::agg_unit(),
215            f,
216            _marker: PhantomData,
217        }
218    }
219}
220
221impl<Spec, L, F> BstSeeker for SeekByRaccCond<Spec, L, F>
222where
223    Spec: BstSpec<Data: BstDataAccess<data::marker::LazyMap, Value = LazyMapElement<L>>>,
224    L: LazyMapMonoid,
225    F: FnMut(&L::Agg) -> bool,
226{
227    type Spec = Spec;
228
229    fn bst_seek(&mut self, node: BstImmutRef<'_, Self::Spec>) -> Ordering {
230        if let Ok(right) = node.reborrow().right().descend() {
231            let right_agg = &right.into_data().bst_data().agg;
232            let nagg = L::agg_operate(right_agg, &self.acc);
233            if (self.f)(&nagg) {
234                return Ordering::Less;
235            }
236            let nagg = L::agg_operate(
237                &L::single_agg(&node.reborrow().into_data().bst_data().key),
238                &nagg,
239            );
240            if (self.f)(&nagg) {
241                Ordering::Equal
242            } else {
243                self.acc = nagg;
244                Ordering::Greater
245            }
246        } else {
247            let nagg = L::agg_operate(
248                &L::single_agg(&node.reborrow().into_data().bst_data().key),
249                &self.acc,
250            );
251            if (self.f)(&nagg) {
252                Ordering::Equal
253            } else {
254                self.acc = nagg;
255                Ordering::Greater
256            }
257        }
258    }
259}