1use super::{
2 Allocator, BoxAllocator, LazyMapMonoid, MonoidAct, Xorshift,
3 binary_search_tree::{
4 BstDataAccess, BstDataMutRef, BstNode, BstNodeId, BstNodeIdManager, BstRoot, BstSeeker,
5 BstSpec,
6 data::{self, LazyMapElement, MonoidActElement},
7 node::WithParent,
8 seeker::{SeekByAccCond, SeekByKey, SeekByRaccCond},
9 split::{Split, Split3},
10 },
11};
12use std::{
13 borrow::Borrow,
14 cmp::Ordering,
15 fmt::{self, Debug},
16 marker::PhantomData,
17 mem::ManuallyDrop,
18 ops::{DerefMut, RangeBounds},
19};
20
21type TreapRoot<M, L> = BstRoot<TreapSpec<M, L>>;
22type TreapNode<M, L> = BstNode<TreapData<M, L>, WithParent<TreapData<M, L>>>;
23
24pub struct TreapSpec<M, L> {
25 _marker: PhantomData<(M, L)>,
26}
27
28pub struct TreapData<M, L>
29where
30 M: MonoidAct<Key: Ord>,
31 L: LazyMapMonoid,
32{
33 priority: u64,
34 key: MonoidActElement<M>,
35 value: LazyMapElement<L>,
36}
37
38impl<M, L> Debug for TreapData<M, L>
39where
40 M: MonoidAct<Key: Ord + Debug, Act: Debug>,
41 L: LazyMapMonoid<Key: Debug, Agg: Debug, Act: Debug>,
42{
43 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
44 f.debug_struct("TreapData")
45 .field("priority", &self.priority)
46 .field("key", &self.key)
47 .field("value", &self.value)
48 .finish()
49 }
50}
51
52impl<M, L> BstDataAccess<data::marker::Key> for TreapData<M, L>
53where
54 M: MonoidAct<Key: Ord>,
55 L: LazyMapMonoid,
56{
57 type Value = M::Key;
58
59 fn bst_data(&self) -> &Self::Value {
60 &self.key.key
61 }
62
63 fn bst_data_mut(&mut self) -> &mut Self::Value {
64 &mut self.key.key
65 }
66}
67
68impl<M, L> BstDataAccess<data::marker::MonoidAct> for TreapData<M, L>
69where
70 M: MonoidAct<Key: Ord>,
71 L: LazyMapMonoid,
72{
73 type Value = MonoidActElement<M>;
74
75 fn bst_data(&self) -> &Self::Value {
76 &self.key
77 }
78
79 fn bst_data_mut(&mut self) -> &mut Self::Value {
80 &mut self.key
81 }
82}
83
84impl<M, L> BstDataAccess<data::marker::LazyMap> for TreapData<M, L>
85where
86 M: MonoidAct<Key: Ord>,
87 L: LazyMapMonoid,
88{
89 type Value = LazyMapElement<L>;
90
91 fn bst_data(&self) -> &Self::Value {
92 &self.value
93 }
94
95 fn bst_data_mut(&mut self) -> &mut Self::Value {
96 &mut self.value
97 }
98}
99
100impl<M, L> BstSpec for TreapSpec<M, L>
101where
102 M: MonoidAct<Key: Ord>,
103 L: LazyMapMonoid,
104{
105 type Parent = WithParent<Self::Data>;
106 type Data = TreapData<M, L>;
107
108 fn top_down(mut node: BstDataMutRef<'_, Self>) {
109 MonoidActElement::<M>::top_down(node.reborrow_datamut());
110 LazyMapElement::<L>::top_down(node.reborrow_datamut());
111 }
112
113 fn bottom_up(mut node: BstDataMutRef<'_, Self>) {
114 LazyMapElement::<L>::bottom_up(node.reborrow_datamut());
115 }
116
117 fn merge(
118 left: Option<TreapRoot<M, L>>,
119 right: Option<TreapRoot<M, L>>,
120 ) -> Option<TreapRoot<M, L>> {
121 match (left, right) {
122 (None, None) => None,
123 (None, Some(node)) | (Some(node), None) => Some(node),
124 (Some(mut left), Some(mut right)) => unsafe {
125 if left.reborrow().into_data().priority > right.reborrow().into_data().priority {
126 TreapSpec::top_down(left.borrow_datamut());
127 let lr = left.borrow_mut().right().take();
128 let lr = Self::merge(lr, Some(right)).unwrap_unchecked();
129 left.borrow_mut().right().set(lr);
130 TreapSpec::bottom_up(left.borrow_datamut());
131 Some(left)
132 } else {
133 TreapSpec::top_down(right.borrow_datamut());
134 let rl = right.borrow_mut().left().take();
135 let rl = Self::merge(Some(left), rl).unwrap_unchecked();
136 right.borrow_mut().left().set(rl);
137 TreapSpec::bottom_up(right.borrow_datamut());
138 Some(right)
139 }
140 },
141 }
142 }
143
144 fn split<Seeker>(
145 node: Option<TreapRoot<M, L>>,
146 mut seeker: Seeker,
147 eq_left: bool,
148 ) -> (Option<TreapRoot<M, L>>, Option<TreapRoot<M, L>>)
149 where
150 Seeker: BstSeeker<Spec = Self>,
151 {
152 match node {
153 None => (None, None),
154 Some(mut node) => {
155 Self::top_down(node.borrow_datamut());
156 let seek_left = match seeker.bst_seek(node.reborrow()) {
157 Ordering::Less => false,
158 Ordering::Equal => !eq_left,
159 Ordering::Greater => true,
160 };
161 if seek_left {
162 unsafe {
163 let left = node.borrow_mut().left().take();
164 let (l, r) = Self::split(left, seeker, eq_left);
165 if let Some(r) = r {
166 node.borrow_mut().left().set(r);
167 }
168 Self::bottom_up(node.borrow_datamut());
169 (l, Some(node))
170 }
171 } else {
172 unsafe {
173 let right = node.borrow_mut().right().take();
174 let (l, r) = Self::split(right, seeker, eq_left);
175 if let Some(l) = l {
176 node.borrow_mut().right().set(l);
177 }
178 Self::bottom_up(node.borrow_datamut());
179 (Some(node), r)
180 }
181 }
182 }
183 }
184 }
185}
186
187impl<M, L> TreapSpec<M, L>
188where
189 M: MonoidAct<Key: Ord>,
190 L: LazyMapMonoid,
191{
192 pub fn merge_ordered(
193 left: Option<TreapRoot<M, L>>,
194 right: Option<TreapRoot<M, L>>,
195 ) -> Option<TreapRoot<M, L>> {
196 match (left, right) {
197 (None, None) => None,
198 (None, Some(node)) | (Some(node), None) => Some(node),
199 (Some(mut left), Some(mut right)) => unsafe {
200 if left.reborrow().into_data().priority > right.reborrow().into_data().priority {
201 Self::top_down(left.borrow_datamut());
202 let key = &left.reborrow().into_data().key.key;
203 let (rl, rr) = Self::split(Some(right), SeekByKey::new(key), false);
204 let ll = left.borrow_mut().left().take();
205 let lr = left.borrow_mut().right().take();
206 if let Some(l) = Self::merge_ordered(ll, rl) {
207 left.borrow_mut().left().set(l);
208 }
209 if let Some(r) = Self::merge_ordered(lr, rr) {
210 left.borrow_mut().right().set(r);
211 }
212 Self::bottom_up(left.borrow_datamut());
213 Some(left)
214 } else {
215 Self::top_down(right.borrow_datamut());
216 let key = &right.reborrow().into_data().key.key;
217 let (ll, lr) = Self::split(Some(left), SeekByKey::new(key), false);
218 let rl = right.borrow_mut().left().take();
219 let rr = right.borrow_mut().right().take();
220 if let Some(l) = Self::merge_ordered(ll, rl) {
221 right.borrow_mut().left().set(l);
222 }
223 if let Some(r) = Self::merge_ordered(lr, rr) {
224 right.borrow_mut().right().set(r);
225 }
226 Self::bottom_up(right.borrow_datamut());
227 Some(right)
228 }
229 },
230 }
231 }
232}
233
234pub struct Treap<M, L, A = BoxAllocator<TreapNode<M, L>>>
235where
236 M: MonoidAct<Key: Ord>,
237 L: LazyMapMonoid,
238 A: Allocator<TreapNode<M, L>>,
239{
240 root: Option<TreapRoot<M, L>>,
241 node_id_manager: BstNodeIdManager<TreapSpec<M, L>>,
242 rng: Xorshift,
243 allocator: ManuallyDrop<A>,
244 _marker: PhantomData<(M, L)>,
245}
246
247impl<M, L, A> Default for Treap<M, L, A>
248where
249 M: MonoidAct<Key: Ord>,
250 L: LazyMapMonoid,
251 A: Allocator<TreapNode<M, L>> + Default,
252{
253 fn default() -> Self {
254 Self {
255 root: None,
256 node_id_manager: Default::default(),
257 rng: Xorshift::new(),
258 allocator: ManuallyDrop::new(A::default()),
259 _marker: PhantomData,
260 }
261 }
262}
263
264impl<M, L, A> Drop for Treap<M, L, A>
265where
266 M: MonoidAct<Key: Ord>,
267 L: LazyMapMonoid,
268 A: Allocator<TreapNode<M, L>>,
269{
270 fn drop(&mut self) {
271 unsafe {
272 if let Some(root) = self.root.take() {
273 root.into_dying().drop_all(self.allocator.deref_mut());
274 }
275 ManuallyDrop::drop(&mut self.allocator);
276 }
277 }
278}
279
280impl<M, L> Treap<M, L>
281where
282 M: MonoidAct<Key: Ord>,
283 L: LazyMapMonoid,
284{
285 pub fn new() -> Self {
286 Self::default()
287 }
288}
289
290impl<M, L, A> Treap<M, L, A>
291where
292 M: MonoidAct<Key: Ord>,
293 L: LazyMapMonoid,
294 A: Allocator<TreapNode<M, L>>,
295{
296 pub fn len(&self) -> usize {
297 self.node_id_manager.len()
298 }
299
300 pub fn is_empty(&self) -> bool {
301 self.node_id_manager.is_empty()
302 }
303
304 pub fn clear(&mut self) {
305 unsafe {
306 if let Some(root) = self.root.take() {
307 root.into_dying().drop_all(self.allocator.deref_mut());
308 }
309 self.node_id_manager.clear();
310 }
311 }
312
313 pub fn get(&mut self, node_id: BstNodeId<TreapSpec<M, L>>) -> Option<(&M::Key, &L::Key)> {
314 if !self.node_id_manager.contains(&node_id) {
315 return None;
316 }
317 unsafe {
318 WithParent::resolve_top_down::<TreapSpec<M, L>>(
319 node_id.reborrow_datamut(&mut self.root),
320 );
321 let data = node_id.reborrow(&self.root).into_data();
322 Some((&data.key.key, &data.value.key))
323 }
324 }
325
326 pub fn change(
327 &mut self,
328 node_id: BstNodeId<TreapSpec<M, L>>,
329 f: impl FnOnce(&mut L::Key),
330 ) -> bool {
331 if !self.node_id_manager.contains(&node_id) {
332 return false;
333 }
334 unsafe {
335 WithParent::resolve_top_down::<TreapSpec<M, L>>(
336 node_id.reborrow_datamut(&mut self.root),
337 );
338 let data = node_id.reborrow_datamut(&mut self.root).into_data_mut();
339 f(&mut data.value.key);
340 WithParent::resolve_bottom_up::<TreapSpec<M, L>>(
341 node_id.reborrow_datamut(&mut self.root),
342 );
343 }
344 true
345 }
346
347 pub fn change_key_value(
348 &mut self,
349 node_id: BstNodeId<TreapSpec<M, L>>,
350 f: impl FnOnce(&mut M::Key, &mut L::Key),
351 ) -> bool {
352 if !self.node_id_manager.contains(&node_id) {
353 return false;
354 }
355 unsafe {
356 WithParent::resolve_top_down::<TreapSpec<M, L>>(
357 node_id.reborrow_datamut(&mut self.root),
358 );
359 let mut node = if WithParent::is_root(node_id.reborrow(&self.root)) {
360 WithParent::remove_root(&mut self.root).unwrap_unchecked()
361 } else {
362 WithParent::remove_not_root(node_id.reborrow_mut(&mut self.root))
363 };
364 let data = node.borrow_datamut().into_data_mut();
365 f(&mut data.key.key, &mut data.value.key);
366 self.root = TreapSpec::merge_ordered(self.root.take(), Some(node));
367 true
368 }
369 }
370
371 pub fn insert(&mut self, key: M::Key, value: L::Key) -> BstNodeId<TreapSpec<M, L>> {
372 let (left, right) = TreapSpec::split(self.root.take(), SeekByKey::new(&key), false);
373 let data = TreapData {
374 priority: self.rng.rand64(),
375 key: MonoidActElement::from_key(key),
376 value: LazyMapElement::from_key(value),
377 };
378 let node = BstRoot::from_data(data, self.allocator.deref_mut());
379 let node_id = self.node_id_manager.register(&node);
380 self.root = TreapSpec::merge(TreapSpec::merge(left, Some(node)), right);
381 node_id
382 }
383
384 pub fn remove(&mut self, node_id: BstNodeId<TreapSpec<M, L>>) -> Option<(M::Key, L::Key)> {
385 if !self.node_id_manager.contains(&node_id) {
386 return None;
387 }
388 unsafe {
389 WithParent::resolve_top_down::<TreapSpec<M, L>>(
390 node_id.reborrow_datamut(&mut self.root),
391 );
392 let node = if WithParent::is_root(node_id.reborrow(&self.root)) {
393 WithParent::remove_root(&mut self.root).unwrap_unchecked()
394 } else {
395 WithParent::remove_not_root(node_id.reborrow_mut(&mut self.root))
396 };
397 self.node_id_manager.unregister(node_id);
398 let data = node.into_dying().into_data(self.allocator.deref_mut());
399 Some((data.key.key, data.value.key))
400 }
401 }
402
403 pub fn range_by_key<Q, R>(&mut self, range: R) -> TreapSplit3<'_, M, L>
404 where
405 M: MonoidAct<Key: Borrow<Q>>,
406 Q: Ord + ?Sized,
407 R: RangeBounds<Q>,
408 {
409 let split3 = Split3::seek_by_key(&mut self.root, range);
410 TreapSplit3 {
411 split3,
412 key_updated: false,
413 }
414 }
415
416 pub fn find_by_key<Q>(&mut self, key: &Q) -> Option<BstNodeId<TreapSpec<M, L>>>
417 where
418 M: MonoidAct<Key: Borrow<Q>>,
419 Q: Ord + ?Sized,
420 {
421 let split = Split::new(
422 &mut self.root,
423 SeekByKey::<TreapSpec<M, L>, M::Key, Q>::new(key),
424 false,
425 );
426 let node = split.right()?.leftmost()?;
427 matches!(node.into_data().key.key.borrow().cmp(key), Ordering::Equal)
428 .then(|| self.node_id_manager.registerd_node_id(node))
429 .flatten()
430 }
431
432 pub fn find_by_acc_cond<F>(&mut self, f: F) -> Option<BstNodeId<TreapSpec<M, L>>>
433 where
434 F: FnMut(&L::Agg) -> bool,
435 {
436 let split = Split::new(
437 &mut self.root,
438 SeekByAccCond::<TreapSpec<M, L>, L, F>::new(f),
439 false,
440 );
441 let node = split.right()?.leftmost()?;
442 self.node_id_manager.registerd_node_id(node)
443 }
444
445 pub fn find_by_racc_cond<F>(&mut self, f: F) -> Option<BstNodeId<TreapSpec<M, L>>>
446 where
447 F: FnMut(&L::Agg) -> bool,
448 {
449 let split = Split::new(
450 &mut self.root,
451 SeekByRaccCond::<TreapSpec<M, L>, L, F>::new(f),
452 true,
453 );
454 let node = split.left()?.rightmost()?;
455 self.node_id_manager.registerd_node_id(node)
456 }
457}
458
459pub struct TreapSplit3<'a, M, L>
460where
461 M: MonoidAct<Key: Ord>,
462 L: LazyMapMonoid,
463{
464 split3: Split3<'a, TreapSpec<M, L>>,
465 key_updated: bool,
466}
467
468impl<'a, M, L> TreapSplit3<'a, M, L>
469where
470 M: MonoidAct<Key: Ord>,
471 L: LazyMapMonoid,
472{
473 pub fn fold(&self) -> L::Agg {
474 if let Some(node) = self.split3.mid() {
475 node.reborrow().into_data().value.agg.clone()
476 } else {
477 L::agg_unit()
478 }
479 }
480
481 pub fn update_key(&mut self, act: M::Act) {
482 if let Some(node) = self.split3.mid_datamut() {
483 MonoidActElement::<M>::update_act(node, &act);
484 self.key_updated = true;
485 }
486 }
487
488 pub fn update_value(&mut self, act: L::Act) {
489 if let Some(node) = self.split3.mid_datamut() {
490 LazyMapElement::<L>::update_act(node, &act);
491 }
492 }
493}
494
495impl<'a, M, L> Drop for TreapSplit3<'a, M, L>
496where
497 M: MonoidAct<Key: Ord>,
498 L: LazyMapMonoid,
499{
500 fn drop(&mut self) {
501 if self.key_updated {
502 self.split3.manually_merge(TreapSpec::merge_ordered);
503 }
504 }
505}
506
507#[cfg(test)]
508mod tests {
509 use super::*;
510 use crate::algebra::{
511 AdditiveOperation, EmptyAct, FlattenAct, RangeMaxRangeAdd, RangeSumRangeAdd,
512 };
513
514 #[test]
515 fn test_treap() {
516 const A: i64 = 100;
517 let mut rng = Xorshift::default();
518 let mut treap = Treap::<FlattenAct<AdditiveOperation<i64>>, RangeMaxRangeAdd<i64>>::new();
519 let mut node_ids = vec![];
520 let mut data = vec![];
521 for _ in 0..10000 {
522 let (l, r) = loop {
523 let l = rng.random(-A..=A);
524 let r = rng.random(-A..=A);
525 if l <= r {
526 break (l, r);
527 }
528 };
529 assert_eq!(data.len(), treap.len());
530 assert_eq!(data.is_empty(), treap.is_empty());
531 match rng.random(0..8) {
532 0 => {
533 let key = rng.random(-A..=A);
534 let value = rng.random(-A..=A);
535 let k = data.partition_point(|(k, _)| *k < key);
536 data.insert(k, (key, value));
537 node_ids.insert(k, treap.insert(key, value));
538 }
539 1 => {
540 if !data.is_empty() {
541 let k = rng.random(0..data.len());
542 let expected = data.remove(k);
543 let result = treap.remove(node_ids.remove(k)).unwrap();
544 assert_eq!(expected, result);
545 }
546 }
547 2 => {
548 let expected: i64 = data
549 .iter()
550 .filter(|(k, _)| (l..r).contains(k))
551 .map(|(_, v)| *v)
552 .max()
553 .unwrap_or(i64::MIN);
554 let result = treap.range_by_key(l..r).fold();
555 assert_eq!(expected, result);
556 }
557 3 => {
558 let add = rng.random(-A..=A);
559 for (k, v) in data.iter_mut() {
560 if (l..r).contains(k) {
561 *v += add;
562 }
563 }
564 treap.range_by_key(l..r).update_value(add);
565 }
566 4 => {
567 let add = rng.random(-A..=A);
568 for (k, _) in data.iter_mut() {
569 if (l..r).contains(k) {
570 *k += add;
571 }
572 }
573 treap.range_by_key(l..r).update_key(add);
574 }
575 5 => {
576 if !data.is_empty() {
577 let k = rng.random(0..data.len());
578 let expected = data[k];
579 let result = treap.get(node_ids[k]).unwrap();
580 assert_eq!(expected, (*result.0, *result.1));
581 }
582 }
583 6 => {
584 if !data.is_empty() {
585 let k = rng.random(0..data.len());
586 let x = rng.random(-A..=A);
587 data[k].1 = x;
588 treap.change(node_ids[k], |value| *value = x);
589 }
590 }
591 _ => {
592 if !data.is_empty() {
593 let k = rng.random(0..data.len());
594 let nk = rng.random(-A..=A);
595 let nv = rng.random(-A..=A);
596 data[k].0 = nk;
597 data[k].1 = nv;
598 treap.change_key_value(node_ids[k], |key, value| {
599 *key = nk;
600 *value = nv;
601 });
602 }
603 }
604 }
605 }
606
607 let mut treap = Treap::<EmptyAct<i64>, RangeSumRangeAdd<i64>>::new();
608 let mut node_ids = vec![];
609 let mut data = vec![];
610 for _ in 0..10000 {
611 let (l, r) = loop {
612 let l = rng.random(-A..=A);
613 let r = rng.random(-A..=A);
614 if l <= r {
615 break (l, r);
616 }
617 };
618 assert_eq!(data.len(), treap.len());
619 assert_eq!(data.is_empty(), treap.is_empty());
620 match rng.random(0..10) {
621 0 => {
622 let key = rng.random(-A..=A);
623 let value = rng.random(1..=A);
624 let k = data.partition_point(|(k, _)| *k < key);
625 data.insert(k, (key, value));
626 node_ids.insert(k, treap.insert(key, value));
627 }
628 1 => {
629 if !data.is_empty() {
630 let k = rng.random(0..data.len());
631 let expected = data.remove(k);
632 let result = treap.remove(node_ids.remove(k)).unwrap();
633 assert_eq!(expected, result);
634 }
635 }
636 2 => {
637 let expected: i64 = data
638 .iter()
639 .filter(|(k, _)| (l..r).contains(k))
640 .map(|(_, v)| *v)
641 .sum();
642 let result = treap.range_by_key(l..r).fold().0;
643 assert_eq!(expected, result);
644 }
645 3 => {
646 let add = rng.random(1..=A);
647 for (k, v) in data.iter_mut() {
648 if (l..r).contains(k) {
649 *v += add;
650 }
651 }
652 treap.range_by_key(l..r).update_value(add);
653 }
654 5 => {
655 if !data.is_empty() {
656 let k = rng.random(0..data.len());
657 let expected = data[k];
658 let result = treap.get(node_ids[k]).unwrap();
659 assert_eq!(expected, (*result.0, *result.1));
660 }
661 }
662 6 => {
663 if !data.is_empty() {
664 let k = rng.random(0..data.len());
665 let x = rng.random(1..=A);
666 data[k].1 = x;
667 treap.change(node_ids[k], |value| *value = x);
668 }
669 }
670 7 => {
671 let key = rng.random(-A..=A);
672 let expected = data.iter().find(|(k, _)| *k == key).cloned();
673 let result = treap.find_by_key(&key).map(|id| treap.get(id).unwrap());
674 assert_eq!(expected, result.map(|(k, v)| (*k, *v)));
675 }
676 8 => {
677 let s = rng.random(0..=A);
678 let mut acc = 0;
679 let expected = data.iter().find_map(|(k, v)| {
680 acc += *v;
681 if acc >= s { Some((*k, *v)) } else { None }
682 });
683 let result = treap
684 .find_by_acc_cond(|agg| agg.0 >= s)
685 .map(|id| treap.get(id).unwrap());
686 assert_eq!(expected, result.map(|(k, v)| (*k, *v)));
687 }
688 _ => {
689 let s = rng.random(0..=A);
690 let mut acc = 0;
691 let expected = data.iter().rev().find_map(|(k, v)| {
692 acc += *v;
693 if acc >= s { Some((*k, *v)) } else { None }
694 });
695 let result = treap
696 .find_by_racc_cond(|agg| agg.0 >= s)
697 .map(|id| treap.get(id).unwrap());
698 assert_eq!(expected, result.map(|(k, v)| (*k, *v)));
699 }
700 }
701 }
702 }
703}