competitive/data_structure/
implicit_treap.rs

1use super::{
2    Allocator, LazyMapMonoid, MemoryPool, Xorshift,
3    binary_search_tree::{
4        BstDataAccess, BstDataMutRef, BstNode, BstRoot, BstSeeker, BstSpec, EqualSide,
5        data::{self, LazyMapElement},
6        node::WithNoParent,
7        seeker::{SeekByAccCond, SeekByRaccCond, SeekBySize},
8        split::{Split, Split3},
9    },
10};
11use std::{
12    fmt::{self, Debug},
13    marker::PhantomData,
14    mem::{ManuallyDrop, replace},
15    ops::{DerefMut, RangeBounds},
16    ptr::NonNull,
17};
18
19type ImplicitTreapRoot<T> = BstRoot<ImplicitTreapSpec<T>>;
20type ImplicitTreapNode<T> = BstNode<ImplicitTreapData<T>>;
21
22pub struct ImplicitTreapSpec<T> {
23    _marker: PhantomData<fn() -> T>,
24}
25
26pub struct ImplicitTreapData<T>
27where
28    T: LazyMapMonoid,
29{
30    priority: u64,
31    value: LazyMapElement<T>,
32    size: usize,
33    rev: bool,
34}
35
36impl<T> Debug for ImplicitTreapData<T>
37where
38    T: LazyMapMonoid<Key: Debug, Agg: Debug, Act: Debug>,
39{
40    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
41        f.debug_struct("ImplicitTreapData")
42            .field("priority", &self.priority)
43            .field("value", &self.value)
44            .field("size", &self.size)
45            .field("rev", &self.rev)
46            .finish()
47    }
48}
49
50impl<T> BstDataAccess<data::marker::Size> for ImplicitTreapData<T>
51where
52    T: LazyMapMonoid,
53{
54    type Value = usize;
55
56    fn bst_data(&self) -> &Self::Value {
57        &self.size
58    }
59
60    fn bst_data_mut(&mut self) -> &mut Self::Value {
61        &mut self.size
62    }
63}
64
65impl<T> BstDataAccess<data::marker::LazyMap> for ImplicitTreapData<T>
66where
67    T: LazyMapMonoid,
68{
69    type Value = LazyMapElement<T>;
70
71    fn bst_data(&self) -> &Self::Value {
72        &self.value
73    }
74
75    fn bst_data_mut(&mut self) -> &mut Self::Value {
76        &mut self.value
77    }
78}
79
80impl<T> ImplicitTreapSpec<T>
81where
82    T: LazyMapMonoid,
83{
84    fn update_act(mut node: BstDataMutRef<'_, Self>, act: &T::Act) {
85        if T::is_act_unit(act) {
86            return;
87        }
88        T::act_operate_assign(&mut node.data_mut().value.act, act);
89        node.data_mut().value.key = T::act_key(&node.reborrow().into_data().value.key, act);
90        if let Some(agg) = T::act_agg(&node.reborrow().into_data().value.agg, act) {
91            node.data_mut().value.agg = agg;
92        } else {
93            Self::top_down(node.reborrow_datamut());
94            Self::bottom_up(node);
95        }
96    }
97
98    fn reverse(mut node: BstDataMutRef<'_, Self>) {
99        node.swap_children();
100        let data = node.data_mut();
101        T::toggle(&mut data.value.agg);
102        data.rev ^= true;
103    }
104}
105
106impl<T> BstSpec for ImplicitTreapSpec<T>
107where
108    T: LazyMapMonoid,
109{
110    type Parent = WithNoParent<Self::Data>;
111    type Data = ImplicitTreapData<T>;
112
113    fn top_down(mut node: BstDataMutRef<'_, Self>) {
114        if !T::is_act_unit(&node.reborrow().into_data().value.act) {
115            let act = replace(&mut node.data_mut().value.act, T::act_unit());
116            if let Ok(left) = node.reborrow_datamut().left().descend() {
117                Self::update_act(left, &act);
118            }
119            if let Ok(right) = node.reborrow_datamut().right().descend() {
120                Self::update_act(right, &act);
121            }
122        }
123        if node.reborrow().into_data().rev {
124            node.data_mut().rev = false;
125            if let Ok(left) = node.reborrow_datamut().left().descend() {
126                Self::reverse(left);
127            }
128            if let Ok(right) = node.reborrow_datamut().right().descend() {
129                Self::reverse(right);
130            }
131        }
132    }
133
134    fn bottom_up(mut node: BstDataMutRef<'_, Self>) {
135        let mut agg = T::single_agg(&node.reborrow().into_data().value.key);
136        let mut size = 1;
137        if let Ok(left) = node.reborrow().left().descend() {
138            let data = left.into_data();
139            agg = T::agg_operate(&data.value.agg, &agg);
140            size += data.size;
141        }
142        if let Ok(right) = node.reborrow().right().descend() {
143            let data = right.into_data();
144            agg = T::agg_operate(&agg, &data.value.agg);
145            size += data.size;
146        }
147        let data = node.data_mut();
148        data.value.agg = agg;
149        data.size = size;
150    }
151
152    fn merge(
153        left: Option<ImplicitTreapRoot<T>>,
154        right: Option<ImplicitTreapRoot<T>>,
155    ) -> Option<ImplicitTreapRoot<T>> {
156        match (left, right) {
157            (None, None) => None,
158            (None, Some(node)) | (Some(node), None) => Some(node),
159            (Some(mut left), Some(mut right)) => unsafe {
160                if left.reborrow().into_data().priority > right.reborrow().into_data().priority {
161                    Self::top_down(left.borrow_datamut());
162                    let lr = left.borrow_mut().right().take();
163                    let lr = Self::merge(lr, Some(right)).unwrap_unchecked();
164                    left.borrow_mut().right().set(lr);
165                    Self::bottom_up(left.borrow_datamut());
166                    Some(left)
167                } else {
168                    Self::top_down(right.borrow_datamut());
169                    let rl = right.borrow_mut().left().take();
170                    let rl = Self::merge(Some(left), rl).unwrap_unchecked();
171                    right.borrow_mut().left().set(rl);
172                    Self::bottom_up(right.borrow_datamut());
173                    Some(right)
174                }
175            },
176        }
177    }
178
179    fn split<Seeker>(
180        node: Option<ImplicitTreapRoot<T>>,
181        mut seeker: Seeker,
182        equal_side: EqualSide,
183    ) -> (Option<ImplicitTreapRoot<T>>, Option<ImplicitTreapRoot<T>>)
184    where
185        Seeker: BstSeeker<Spec = Self>,
186    {
187        match node {
188            None => (None, None),
189            Some(mut node) => {
190                Self::top_down(node.borrow_datamut());
191                if equal_side.goes_left(seeker.bst_seek(node.reborrow())) {
192                    unsafe {
193                        let right = node.borrow_mut().right().take();
194                        let (l, r) = Self::split(right, seeker, equal_side);
195                        if let Some(l) = l {
196                            node.borrow_mut().right().set(l);
197                        }
198                        Self::bottom_up(node.borrow_datamut());
199                        (Some(node), r)
200                    }
201                } else {
202                    unsafe {
203                        let left = node.borrow_mut().left().take();
204                        let (l, r) = Self::split(left, seeker, equal_side);
205                        if let Some(r) = r {
206                            node.borrow_mut().left().set(r);
207                        }
208                        Self::bottom_up(node.borrow_datamut());
209                        (l, Some(node))
210                    }
211                }
212            }
213        }
214    }
215}
216
217pub struct ImplicitTreap<T, A = MemoryPool<ImplicitTreapNode<T>>>
218where
219    T: LazyMapMonoid,
220    A: Allocator<ImplicitTreapNode<T>>,
221{
222    root: Option<ImplicitTreapRoot<T>>,
223    length: usize,
224    rng: Xorshift,
225    allocator: ManuallyDrop<A>,
226    _marker: PhantomData<fn() -> T>,
227}
228
229impl<T, A> Default for ImplicitTreap<T, A>
230where
231    T: LazyMapMonoid,
232    A: Allocator<ImplicitTreapNode<T>> + Default,
233{
234    fn default() -> Self {
235        Self {
236            root: None,
237            length: 0,
238            rng: Xorshift::new(),
239            allocator: ManuallyDrop::new(A::default()),
240            _marker: PhantomData,
241        }
242    }
243}
244
245impl<T, A> Drop for ImplicitTreap<T, A>
246where
247    T: LazyMapMonoid,
248    A: Allocator<ImplicitTreapNode<T>>,
249{
250    fn drop(&mut self) {
251        unsafe {
252            if let Some(root) = self.root.take() {
253                root.into_dying().drop_all(self.allocator.deref_mut());
254            }
255            ManuallyDrop::drop(&mut self.allocator);
256        }
257    }
258}
259
260impl<T> ImplicitTreap<T>
261where
262    T: LazyMapMonoid,
263{
264    pub fn new() -> Self {
265        Self::default()
266    }
267
268    pub fn with_capacity(capacity: usize) -> Self {
269        Self {
270            root: None,
271            length: 0,
272            rng: Xorshift::new(),
273            allocator: ManuallyDrop::new(MemoryPool::with_capacity(capacity)),
274            _marker: PhantomData,
275        }
276    }
277}
278
279impl<T, A> ImplicitTreap<T, A>
280where
281    T: LazyMapMonoid,
282    A: Allocator<ImplicitTreapNode<T>>,
283{
284    fn node(&mut self, key: T::Key) -> ImplicitTreapRoot<T> {
285        BstRoot::from_data(
286            ImplicitTreapData {
287                priority: self.rng.rand64(),
288                value: LazyMapElement::from_key(key),
289                size: 1,
290                rev: false,
291            },
292            self.allocator.deref_mut(),
293        )
294    }
295
296    fn build<I>(&mut self, iter: I) -> (Option<ImplicitTreapRoot<T>>, usize)
297    where
298        I: IntoIterator<Item = T::Key>,
299    {
300        let mut stack = vec![];
301        let mut len = 0;
302        for key in iter {
303            let mut cur = self.node(key).node;
304            let mut left = None;
305            unsafe {
306                while stack
307                    .last()
308                    .is_some_and(|node: &NonNull<ImplicitTreapNode<T>>| {
309                        node.as_ref().data.priority < cur.as_ref().data.priority
310                    })
311                {
312                    left = stack.pop();
313                }
314                cur.as_mut().child[0] = left;
315                if let Some(parent) = stack.last_mut() {
316                    parent.as_mut().child[1] = Some(cur);
317                }
318            }
319            stack.push(cur);
320            len += 1;
321        }
322        let root = stack.first().copied().map(BstRoot::new);
323        if let Some(mut root) = root {
324            Self::build_bottom_up(root.borrow_datamut());
325            (Some(root), len)
326        } else {
327            (None, len)
328        }
329    }
330
331    fn build_bottom_up(mut node: BstDataMutRef<'_, ImplicitTreapSpec<T>>) {
332        if let Ok(left) = node.reborrow_datamut().left().descend() {
333            Self::build_bottom_up(left);
334        }
335        if let Ok(right) = node.reborrow_datamut().right().descend() {
336            Self::build_bottom_up(right);
337        }
338        ImplicitTreapSpec::<T>::bottom_up(node);
339    }
340
341    pub fn len(&self) -> usize {
342        self.length
343    }
344
345    pub fn is_empty(&self) -> bool {
346        self.length == 0
347    }
348
349    pub fn update<R>(&mut self, range: R, x: T::Act)
350    where
351        R: RangeBounds<usize>,
352    {
353        let mut split = Split3::seek_by_size(&mut self.root, range);
354        if let Some(root) = split.mid_datamut() {
355            ImplicitTreapSpec::<T>::update_act(root, &x);
356        }
357    }
358
359    pub fn fold<R>(&mut self, range: R) -> T::Agg
360    where
361        R: RangeBounds<usize>,
362    {
363        let split = Split3::seek_by_size(&mut self.root, range);
364        split
365            .mid()
366            .map(|node| node.into_data().value.agg.clone())
367            .unwrap_or_else(T::agg_unit)
368    }
369
370    pub fn reverse<R>(&mut self, range: R)
371    where
372        R: RangeBounds<usize>,
373    {
374        let mut split = Split3::seek_by_size(&mut self.root, range);
375        if let Some(root) = split.mid_datamut() {
376            ImplicitTreapSpec::<T>::reverse(root);
377        }
378    }
379
380    pub fn get(&mut self, index: usize) -> Option<&T::Key> {
381        if index >= self.length {
382            return None;
383        }
384        let split = Split3::seek_by_size(&mut self.root, index..=index);
385        let node = split.mid()?.node;
386        drop(split);
387        Some(unsafe { &(*node.as_ptr()).data.value.key })
388    }
389
390    pub fn modify<F>(&mut self, index: usize, f: F)
391    where
392        F: FnOnce(&T::Key) -> T::Key,
393    {
394        assert!(index < self.length);
395        let mut split = Split3::seek_by_size(&mut self.root, index..=index);
396        let mut node = split.mid_datamut().unwrap();
397        ImplicitTreapSpec::<T>::top_down(node.reborrow_datamut());
398        {
399            let data = node.data_mut();
400            data.value.key = f(&data.value.key);
401        }
402        ImplicitTreapSpec::<T>::bottom_up(node);
403    }
404
405    pub fn insert(&mut self, index: usize, x: T::Key) {
406        assert!(index <= self.length);
407        let node = self.node(x);
408        if index == 0 {
409            self.root = ImplicitTreapSpec::<T>::merge(Some(node), self.root.take());
410        } else if index == self.length {
411            self.root = ImplicitTreapSpec::<T>::merge(self.root.take(), Some(node));
412        } else {
413            let mut node = Some(node);
414            let mut split = Split::new(&mut self.root, SeekBySize::new(index), EqualSide::Right);
415            split.manually_merge(|left, right| {
416                ImplicitTreapSpec::<T>::merge(
417                    ImplicitTreapSpec::<T>::merge(left, node.take()),
418                    right,
419                )
420            });
421        }
422        self.length += 1;
423    }
424
425    pub fn remove(&mut self, index: usize) -> Option<T::Key> {
426        if index >= self.length {
427            return None;
428        }
429        let mid;
430        if index == 0 {
431            let (left, right) = ImplicitTreapSpec::<T>::split(
432                self.root.take(),
433                SeekBySize::new(1),
434                EqualSide::Right,
435            );
436            mid = left;
437            self.root = right;
438        } else if index + 1 == self.length {
439            let (left, right) = ImplicitTreapSpec::<T>::split(
440                self.root.take(),
441                SeekBySize::new(index),
442                EqualSide::Right,
443            );
444            mid = right;
445            self.root = left;
446        } else {
447            let (left, rest) = ImplicitTreapSpec::<T>::split(
448                self.root.take(),
449                SeekBySize::new(index),
450                EqualSide::Right,
451            );
452            let (middle, right) =
453                ImplicitTreapSpec::<T>::split(rest, SeekBySize::new(1), EqualSide::Right);
454            mid = middle;
455            self.root = ImplicitTreapSpec::<T>::merge(left, right);
456        }
457        self.length -= 1;
458        let mut node = mid.unwrap();
459        ImplicitTreapSpec::<T>::top_down(node.borrow_datamut());
460        let data = unsafe { node.into_dying().into_data(self.allocator.deref_mut()) };
461        Some(data.value.key)
462    }
463
464    pub fn partition_point_acc<F>(&mut self, left: usize, mut pred: F) -> usize
465    where
466        F: FnMut(&T::Agg) -> bool,
467    {
468        let mut split3 = Split3::seek_by_size(&mut self.root, left..);
469        let front_size = split3
470            .left()
471            .map(|node| node.into_data().size)
472            .unwrap_or_default();
473        let split = split3.split_mid(SeekByAccCond::new(|acc| !pred(acc)), EqualSide::Right);
474        let index = split
475            .left()
476            .map(|node| node.into_data().size)
477            .unwrap_or_default();
478        front_size + index
479    }
480
481    pub fn rpartition_point_acc<F>(&mut self, right: usize, mut pred: F) -> usize
482    where
483        F: FnMut(&T::Agg) -> bool,
484    {
485        let mut split3 = Split3::seek_by_size(&mut self.root, ..right);
486        let split = split3.split_mid(SeekByRaccCond::new(|acc| !pred(acc)), EqualSide::Left);
487        split
488            .left()
489            .map(|node| node.into_data().size)
490            .unwrap_or_default()
491    }
492
493    pub fn rotate_left(&mut self, mid: usize) {
494        assert!(mid <= self.length);
495        if mid == 0 || mid == self.length {
496            return;
497        }
498        let (left, right) =
499            ImplicitTreapSpec::<T>::split(self.root.take(), SeekBySize::new(mid), EqualSide::Right);
500        self.root = ImplicitTreapSpec::<T>::merge(right, left);
501    }
502
503    pub fn rotate_right(&mut self, k: usize) {
504        assert!(k <= self.length);
505        self.rotate_left(self.length - k);
506    }
507}
508
509impl<T, A> Extend<T::Key> for ImplicitTreap<T, A>
510where
511    T: LazyMapMonoid,
512    A: Allocator<ImplicitTreapNode<T>>,
513{
514    fn extend<I>(&mut self, iter: I)
515    where
516        I: IntoIterator<Item = T::Key>,
517    {
518        let (root, len) = self.build(iter);
519        self.root = ImplicitTreapSpec::<T>::merge(self.root.take(), root);
520        self.length += len;
521    }
522}
523
524#[cfg(test)]
525mod tests {
526    use super::*;
527    use crate::{
528        algebra::{RangeChminChmaxAdd, RangeSumRangeChminChmaxAdd},
529        num::Saturating,
530        tools::{NotEmptySegment, Xorshift},
531    };
532
533    #[test]
534    fn test_implicit_treap_range_sum_chmin_chmax_add_random() {
535        const N: usize = 1_000;
536        const Q: usize = 20_000;
537        const A: i64 = 1_000;
538
539        let mut rng = Xorshift::default();
540        let mut arr: Vec<_> = (0..N).map(|_| Saturating(rng.random(0..=A))).collect();
541        let mut treap = ImplicitTreap::<RangeSumRangeChminChmaxAdd<_>>::with_capacity(N + Q);
542        treap.extend(arr.iter().copied());
543
544        assert_eq!(
545            Saturating(0),
546            ImplicitTreap::<RangeSumRangeChminChmaxAdd<Saturating<i64>>>::new()
547                .fold(..)
548                .sum
549        );
550        assert_eq!(None, treap.remove(N));
551        assert_eq!(None, treap.get(N));
552
553        for _ in 0..Q {
554            assert_eq!(arr.len(), treap.len());
555            assert_eq!(arr.is_empty(), treap.is_empty());
556            match rng.random(0..10) {
557                0 if arr.len() < N * 2 => {
558                    let i = rng.random(0..=arr.len());
559                    let x = Saturating(rng.random(0..=A));
560                    treap.insert(i, x);
561                    arr.insert(i, x);
562                }
563                1 if !arr.is_empty() => {
564                    let i = rng.random(0..arr.len());
565                    assert_eq!(arr.remove(i), treap.remove(i).unwrap());
566                }
567                2 if !arr.is_empty() => {
568                    let (l, r) = rng.random(NotEmptySegment(arr.len()));
569                    assert_eq!(
570                        arr[l..r].iter().copied().sum::<Saturating<i64>>(),
571                        treap.fold(l..r).sum
572                    );
573                }
574                3 if !arr.is_empty() => {
575                    let (l, r) = rng.random(NotEmptySegment(arr.len()));
576                    match rng.random(0..3) {
577                        0 => {
578                            let x = Saturating(rng.random(0..=A));
579                            treap.update(l..r, RangeChminChmaxAdd::chmin(x));
580                            arr[l..r].iter_mut().for_each(|a| *a = (*a).min(x));
581                        }
582                        1 => {
583                            let x = Saturating(rng.random(0..=A));
584                            treap.update(l..r, RangeChminChmaxAdd::chmax(x));
585                            arr[l..r].iter_mut().for_each(|a| *a = (*a).max(x));
586                        }
587                        _ => {
588                            let x = Saturating(rng.random(0..=A));
589                            treap.update(l..r, RangeChminChmaxAdd::add(x));
590                            arr[l..r].iter_mut().for_each(|a| *a += x);
591                        }
592                    }
593                }
594                4 if !arr.is_empty() => {
595                    let (l, r) = rng.random(NotEmptySegment(arr.len()));
596                    treap.reverse(l..r);
597                    arr[l..r].reverse();
598                }
599                5 if !arr.is_empty() => {
600                    let left = rng.random(0..=arr.len());
601                    let sum = arr[left..].iter().copied().sum::<Saturating<i64>>();
602                    let x = Saturating(rng.random(1..=sum.0.saturating_add(A)));
603                    assert_eq!(
604                        treap.partition_point_acc(left, |acc| acc.sum < x),
605                        arr[left..]
606                            .iter()
607                            .scan(Saturating(0), |acc, &a| {
608                                *acc += a;
609                                Some(*acc)
610                            })
611                            .position(|acc| acc >= x)
612                            .map_or(arr.len(), |i| i + left),
613                    );
614                }
615                6 if !arr.is_empty() => {
616                    let right = rng.random(0..=arr.len());
617                    let sum = arr[..right].iter().copied().sum::<Saturating<i64>>();
618                    let x = Saturating(rng.random(1..=sum.0.saturating_add(A)));
619                    assert_eq!(
620                        treap.rpartition_point_acc(right, |acc| acc.sum < x),
621                        arr[..right]
622                            .iter()
623                            .rev()
624                            .scan(Saturating(0), |acc, &a| {
625                                *acc += a;
626                                Some(*acc)
627                            })
628                            .position(|acc| acc >= x)
629                            .map_or(0, |i| right - i),
630                    );
631                }
632                7 => {
633                    let i = rng.random(0..=arr.len());
634                    treap.rotate_left(i);
635                    arr.rotate_left(i);
636                }
637                8 => {
638                    let i = rng.random(0..=arr.len());
639                    treap.rotate_right(i);
640                    arr.rotate_right(i);
641                }
642                _ if !arr.is_empty() => {
643                    let i = rng.random(0..arr.len());
644                    if rng.random(0..2) == 0 {
645                        assert_eq!(arr.get(i), treap.get(i));
646                    } else {
647                        let x = Saturating(rng.random(0..=A));
648                        treap.modify(i, |_| x);
649                        arr[i] = x;
650                    }
651                }
652                _ => {}
653            }
654        }
655    }
656}
competitive/data_structure/implicit_treap.rs

competitive/data_structure/
implicit_treap.rs
