competitive/data_structure/binary_search_tree/
seeker.rs1use 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}