competitive/data_structure/splay_tree/
sequence.rs

1use super::{
2    Allocator, LazyMapMonoid, MemoryPool,
3    node::{Node, NodeRange, NodeRef, Root, SplaySeeker, SplaySpec, marker},
4};
5use std::{
6    cmp::Ordering,
7    fmt::{self, Debug},
8    marker::PhantomData,
9    mem::{ManuallyDrop, replace},
10    ops::{Bound, DerefMut, RangeBounds},
11};
12
13pub struct LazyAggElement<T>
14where
15    T: LazyMapMonoid,
16{
17    key: T::Key,
18    agg: T::Agg,
19    lazy: T::Act,
20    size: usize,
21    rev: bool,
22}
23
24impl<T> Debug for LazyAggElement<T>
25where
26    T: LazyMapMonoid<Key: Debug, Agg: Debug, Act: Debug>,
27{
28    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
29        f.debug_struct("LazyAggElement")
30            .field("key", &self.key)
31            .field("agg", &self.agg)
32            .field("lazy", &self.lazy)
33            .field("size", &self.size)
34            .finish()
35    }
36}
37
38pub struct LazyAggSplay<T> {
39    _marker: PhantomData<fn() -> T>,
40}
41
42impl<T> LazyAggSplay<T>
43where
44    T: LazyMapMonoid,
45{
46    pub fn update_lazy(mut node: NodeRef<marker::DataMut<'_>, Self>, lazy: &T::Act) {
47        T::act_operate_assign(&mut node.data_mut().lazy, lazy);
48        node.data_mut().key = T::act_key(&node.data().key, lazy);
49        if let Some(nxlazy) = T::act_agg(&node.data().agg, lazy) {
50            node.data_mut().agg = nxlazy;
51        } else {
52            node = Self::propagate(node);
53            Self::recalc(node);
54        }
55    }
56    pub fn reverse(node: NodeRef<marker::DataMut<'_>, Self>) {
57        node.reverse();
58        T::toggle(&mut node.data_mut().agg);
59        node.data_mut().rev ^= true;
60    }
61    fn propagate(node: NodeRef<marker::DataMut<'_>, Self>) -> NodeRef<marker::DataMut<'_>, Self> {
62        let lazy = replace(&mut node.data_mut().lazy, T::act_unit());
63        if let Some(left) = node.left() {
64            Self::update_lazy(left, &lazy);
65        }
66        if let Some(right) = node.right() {
67            Self::update_lazy(right, &lazy);
68        }
69        if replace(&mut node.data_mut().rev, false) {
70            if let Some(left) = node.left() {
71                Self::reverse(left);
72            }
73            if let Some(right) = node.right() {
74                Self::reverse(right);
75            }
76        }
77        node
78    }
79    fn recalc(node: NodeRef<marker::DataMut<'_>, Self>) -> NodeRef<marker::DataMut<'_>, Self> {
80        let mut agg = T::single_agg(&node.data().key);
81        let mut size = 1;
82        if let Some(left) = node.left() {
83            let data = left.data();
84            agg = T::agg_operate(&data.agg, &agg);
85            // agg = <T::AggMonoid as Magma>::operate(&data.agg, &agg);
86            size += data.size;
87        }
88        if let Some(right) = node.right() {
89            let data = right.data();
90            agg = T::agg_operate(&agg, &data.agg);
91            size += data.size;
92        }
93        let data = node.data_mut();
94        data.agg = agg;
95        data.size = size;
96        node
97    }
98}
99impl<T> SplaySpec for LazyAggSplay<T>
100where
101    T: LazyMapMonoid,
102{
103    type T = LazyAggElement<T>;
104    fn has_bottom_up() -> bool {
105        true
106    }
107    fn top_down(node: NodeRef<marker::DataMut<'_>, Self>) {
108        Self::propagate(node);
109    }
110    fn bottom_up(node: NodeRef<marker::DataMut<'_>, Self>) {
111        Self::recalc(node);
112    }
113}
114
115struct SeekBySize<T> {
116    index: usize,
117    _marker: PhantomData<fn() -> T>,
118}
119impl<T> SeekBySize<T> {
120    fn new(index: usize) -> Self {
121        Self {
122            index,
123            _marker: PhantomData,
124        }
125    }
126}
127impl<T> SplaySeeker for SeekBySize<T>
128where
129    T: LazyMapMonoid,
130{
131    type S = LazyAggSplay<T>;
132    fn splay_seek(&mut self, node: NodeRef<marker::Immut<'_>, Self::S>) -> Ordering {
133        let lsize = node.left().map(|l| l.data().size).unwrap_or_default();
134        let ord = self.index.cmp(&lsize);
135        if matches!(ord, Ordering::Greater) {
136            self.index -= lsize + 1;
137        }
138        ord
139    }
140}
141
142struct SeekByAccCond<F, T>
143where
144    T: LazyMapMonoid,
145{
146    acc: T::Agg,
147    f: F,
148    _marker: PhantomData<fn() -> T>,
149}
150impl<F, T> SeekByAccCond<F, T>
151where
152    T: LazyMapMonoid,
153{
154    fn new(f: F) -> Self {
155        Self {
156            acc: T::agg_unit(),
157            f,
158            _marker: PhantomData,
159        }
160    }
161}
162impl<F, T> SplaySeeker for SeekByAccCond<F, T>
163where
164    F: FnMut(&T::Agg) -> bool,
165    T: LazyMapMonoid,
166{
167    type S = LazyAggSplay<T>;
168    fn splay_seek(&mut self, node: NodeRef<marker::Immut<'_>, Self::S>) -> Ordering {
169        if let Some(lagg) = node.left().map(|l| &l.data().agg) {
170            let nacc = T::agg_operate(&self.acc, lagg);
171            if (self.f)(&nacc) {
172                return Ordering::Less;
173            }
174            self.acc = nacc;
175        };
176        self.acc = T::agg_operate(&self.acc, &T::single_agg(&node.data().key));
177        if (self.f)(&self.acc) {
178            Ordering::Equal
179        } else {
180            Ordering::Greater
181        }
182    }
183}
184
185struct SeekByRaccCond<F, T>
186where
187    T: LazyMapMonoid,
188{
189    acc: T::Agg,
190    f: F,
191    _marker: PhantomData<fn() -> T>,
192}
193impl<F, T> SeekByRaccCond<F, T>
194where
195    T: LazyMapMonoid,
196{
197    fn new(f: F) -> Self {
198        Self {
199            acc: T::agg_unit(),
200            f,
201            _marker: PhantomData,
202        }
203    }
204}
205impl<F, T> SplaySeeker for SeekByRaccCond<F, T>
206where
207    F: FnMut(&T::Agg) -> bool,
208    T: LazyMapMonoid,
209{
210    type S = LazyAggSplay<T>;
211    fn splay_seek(&mut self, node: NodeRef<marker::Immut<'_>, Self::S>) -> Ordering {
212        if let Some(lagg) = node.right().map(|r| &r.data().agg) {
213            let nacc = T::agg_operate(lagg, &self.acc);
214            if (self.f)(&nacc) {
215                return Ordering::Greater;
216            }
217            self.acc = nacc;
218        };
219        self.acc = T::agg_operate(&T::single_agg(&node.data().key), &self.acc);
220        if (self.f)(&self.acc) {
221            Ordering::Equal
222        } else {
223            Ordering::Less
224        }
225    }
226}
227
228pub struct SplaySequence<T, A = MemoryPool<Node<LazyAggElement<T>>>>
229where
230    T: LazyMapMonoid,
231    A: Allocator<Node<LazyAggElement<T>>>,
232{
233    root: Root<LazyAggSplay<T>>,
234    length: usize,
235    alloc: ManuallyDrop<A>,
236}
237
238impl<T, A> Debug for SplaySequence<T, A>
239where
240    T: LazyMapMonoid<Key: Debug, Agg: Debug, Act: Debug>,
241    A: Allocator<Node<LazyAggElement<T>>>,
242{
243    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
244        f.debug_struct("SplayMap")
245            .field("root", &self.root)
246            .field("length", &self.length)
247            .finish()
248    }
249}
250
251impl<T, A> Drop for SplaySequence<T, A>
252where
253    T: LazyMapMonoid,
254    A: Allocator<Node<LazyAggElement<T>>>,
255{
256    fn drop(&mut self) {
257        unsafe {
258            while let Some(node) = self.root.take_first() {
259                self.alloc.deallocate(node.into_dying().into_inner());
260            }
261            ManuallyDrop::drop(&mut self.alloc);
262        }
263    }
264}
265
266impl<T, A> Default for SplaySequence<T, A>
267where
268    T: LazyMapMonoid,
269    A: Allocator<Node<LazyAggElement<T>>> + Default,
270{
271    fn default() -> Self {
272        Self {
273            root: Root::default(),
274            length: 0,
275            alloc: Default::default(),
276        }
277    }
278}
279
280impl<T> SplaySequence<T>
281where
282    T: LazyMapMonoid,
283{
284    pub fn new() -> Self {
285        Default::default()
286    }
287    pub fn with_capacity(capacity: usize) -> Self {
288        Self {
289            root: Root::default(),
290            length: 0,
291            alloc: ManuallyDrop::new(MemoryPool::with_capacity(capacity)),
292        }
293    }
294    pub fn len(&self) -> usize {
295        self.length
296    }
297    pub fn is_empty(&self) -> bool {
298        self.length == 0
299    }
300}
301impl<T, A> SplaySequence<T, A>
302where
303    T: LazyMapMonoid,
304    A: Allocator<Node<LazyAggElement<T>>>,
305{
306    fn range<R>(&mut self, range: R) -> NodeRange<'_, LazyAggSplay<T>>
307    where
308        R: RangeBounds<usize>,
309    {
310        let start = match range.start_bound() {
311            Bound::Included(&index) => Bound::Included(SeekBySize::new(index)),
312            Bound::Excluded(&index) => Bound::Excluded(SeekBySize::new(index)),
313            Bound::Unbounded => Bound::Unbounded,
314        };
315        let end = match range.end_bound() {
316            Bound::Included(&index) => Bound::Included(SeekBySize::new(index)),
317            Bound::Excluded(&index) => Bound::Excluded(SeekBySize::new(index)),
318            Bound::Unbounded => Bound::Unbounded,
319        };
320        self.root.range(start, end)
321    }
322    pub fn update<R>(&mut self, range: R, x: T::Act)
323    where
324        R: RangeBounds<usize>,
325    {
326        if let Some(root) = self.range(range).root_mut().root_data_mut() {
327            LazyAggSplay::<T>::update_lazy(root, &x);
328        }
329    }
330    pub fn fold<R>(&mut self, range: R) -> T::Agg
331    where
332        R: RangeBounds<usize>,
333    {
334        if let Some(root) = self.range(range).root().root() {
335            root.data().agg.clone()
336        } else {
337            T::agg_unit()
338        }
339    }
340    pub fn reverse<R>(&mut self, range: R)
341    where
342        R: RangeBounds<usize>,
343    {
344        if let Some(root) = self.range(range).root_mut().root_data_mut() {
345            LazyAggSplay::<T>::reverse(root);
346        }
347    }
348    pub fn get(&mut self, index: usize) -> Option<&T::Key> {
349        self.root.splay_by(SeekBySize::new(index))?;
350        self.root.root().map(|root| &root.data().key)
351    }
352    pub fn modify<F>(&mut self, index: usize, f: F)
353    where
354        F: FnOnce(&T::Key) -> T::Key,
355    {
356        self.root.splay_by(SeekBySize::new(index)).unwrap();
357        let data = self.root.root_data_mut().unwrap().data_mut();
358        data.key = f(&data.key);
359        LazyAggSplay::<T>::bottom_up(self.root.root_data_mut().unwrap());
360    }
361    pub fn insert(&mut self, index: usize, x: T::Key) {
362        assert!(index <= self.length);
363        self.root.splay_by(SeekBySize::new(index));
364        let agg = T::single_agg(&x);
365        unsafe {
366            let node = NodeRef::from_data(
367                LazyAggElement {
368                    key: x,
369                    agg,
370                    lazy: T::act_unit(),
371                    size: 1,
372                    rev: false,
373                },
374                self.alloc.deref_mut(),
375            );
376            if index == self.length {
377                self.root.insert_right(node);
378            } else {
379                self.root.insert_left(node);
380            }
381        }
382        self.length += 1;
383    }
384    pub fn remove(&mut self, index: usize) -> Option<T::Key> {
385        if index >= self.length {
386            return None;
387        }
388        self.root.splay_by(SeekBySize::new(index));
389        self.length -= 1;
390        let node = self.root.take_root().unwrap().into_dying();
391        unsafe { Some(node.into_data(self.alloc.deref_mut()).key) }
392    }
393    pub fn position_acc<R, F>(&mut self, range: R, f: F) -> Option<usize>
394    where
395        R: RangeBounds<usize>,
396        F: FnMut(&T::Agg) -> bool,
397    {
398        let mut range = self.range(range);
399        let ord = range.root_mut().splay_by(SeekByAccCond::new(f));
400        if !matches!(ord, Some(Ordering::Equal)) {
401            return None;
402        }
403        let front_size = range.front().size();
404        let left_size = range.root().left_size();
405        Some(front_size + left_size)
406    }
407    pub fn rposition_acc<R, F>(&mut self, range: R, f: F) -> Option<usize>
408    where
409        R: RangeBounds<usize>,
410        F: FnMut(&T::Agg) -> bool,
411    {
412        let mut range = self.range(range);
413        let ord = range.root_mut().splay_by(SeekByRaccCond::new(f));
414        if !matches!(ord, Some(Ordering::Equal)) {
415            return None;
416        }
417        let front_size = range.front().size();
418        let left_size = range.root().left_size();
419        Some(front_size + left_size)
420    }
421    pub fn rotate_left(&mut self, mid: usize) {
422        assert!(mid <= self.length);
423        if mid != 0 || mid != self.length {
424            self.range(mid..).drop_rotate_left()
425        }
426    }
427    pub fn rotate_right(&mut self, k: usize) {
428        assert!(k <= self.length);
429        self.rotate_left(self.length - k);
430    }
431}
432
433impl<T, A> Extend<T::Key> for SplaySequence<T, A>
434where
435    T: LazyMapMonoid,
436    A: Allocator<Node<LazyAggElement<T>>>,
437{
438    fn extend<I>(&mut self, iter: I)
439    where
440        I: IntoIterator<Item = T::Key>,
441    {
442        let nodes: Vec<_> = unsafe {
443            iter.into_iter()
444                .map(|key| {
445                    let agg = T::single_agg(&key);
446                    NodeRef::from_data(
447                        LazyAggElement {
448                            key,
449                            agg,
450                            lazy: T::act_unit(),
451                            size: 1,
452                            rev: false,
453                        },
454                        self.alloc.deref_mut(),
455                    )
456                })
457                .collect()
458        };
459        self.length += nodes.len();
460        let mut root = unsafe { Root::from_single_nodes(nodes) };
461        self.root.append(&mut root);
462    }
463}
464
465impl<T> Root<LazyAggSplay<T>>
466where
467    T: LazyMapMonoid,
468{
469    fn size(&self) -> usize {
470        self.root().map(|root| root.data().size).unwrap_or_default()
471    }
472    fn left_size(&self) -> usize {
473        self.root()
474            .and_then(|root| root.left().map(|left| left.data().size))
475            .unwrap_or_default()
476    }
477}
478
479#[cfg(test)]
480mod tests {
481    use super::*;
482    use crate::{
483        algebra::RangeMaxRangeUpdate,
484        rand,
485        tools::{NotEmptySegment, Xorshift},
486    };
487
488    #[test]
489    fn test_splay_sequence() {
490        const N: usize = 1_000;
491        const Q: usize = 20_000;
492        const A: i64 = 1_000_000_000;
493
494        let mut rng = Xorshift::default();
495        rand!(rng, mut arr: [-A..A; N]);
496        let mut seq = SplaySequence::<RangeMaxRangeUpdate<_>>::new();
497        seq.extend(arr.iter().cloned());
498        for _ in 0..Q {
499            assert_eq!(arr.len(), seq.len());
500            rand!(rng, ty: 0..6, (l, r): NotEmptySegment(arr.len()));
501            match ty {
502                0 if arr.len() < N * 2 => {
503                    rand!(rng, i: ..=arr.len(), x: -A..A);
504                    seq.insert(i, x);
505                    arr.insert(i, x);
506                }
507                1 if arr.len() > 1 => {
508                    rand!(rng, i: ..arr.len());
509                    assert_eq!(arr.remove(i), seq.remove(i).unwrap());
510                }
511                2 => {
512                    let res = arr[l..r].iter().max().cloned().unwrap_or_default();
513                    assert_eq!(seq.fold(l..r), res);
514                }
515                3 => {
516                    rand!(rng, x: -A..A);
517                    seq.update(l..r, Some(x));
518                    arr[l..r].iter_mut().for_each(|a| *a = x);
519                }
520                4 => {
521                    arr[l..r].reverse();
522                    seq.reverse(l..r);
523                }
524                5 => {
525                    rand!(rng, x: -A..A);
526                    assert_eq!(
527                        seq.position_acc(l..r, |&d| d >= x),
528                        arr[l..r]
529                            .iter()
530                            .scan(i64::MIN, |acc, &a| {
531                                *acc = a.max(*acc);
532                                Some(*acc)
533                            })
534                            .position(|acc| acc >= x)
535                            .map(|i| i + l),
536                    );
537                }
538                6 => {
539                    rand!(rng, x: -A..A);
540                    assert_eq!(
541                        seq.rposition_acc(l..r, |&d| d >= x),
542                        arr[l..r]
543                            .iter()
544                            .rev()
545                            .scan(i64::MIN, |acc, &a| {
546                                *acc = a.max(*acc);
547                                Some(*acc)
548                            })
549                            .position(|acc| acc >= x)
550                            .map(|i| r - i - 1),
551                    );
552                }
553                7 => {
554                    rand!(rng, i: ..=arr.len());
555                    seq.rotate_left(i);
556                    arr.rotate_left(i);
557                }
558                8 => {
559                    rand!(rng, i: ..=arr.len());
560                    seq.rotate_right(i);
561                    arr.rotate_right(i);
562                }
563                _ => {
564                    rand!(rng, i: ..arr.len());
565                    assert_eq!(arr.get(i), seq.get(i));
566                }
567            }
568        }
569    }
570}
competitive/data_structure/splay_tree/sequence.rs

competitive/data_structure/splay_tree/
sequence.rs
