competitive/data_structure/splay_tree/
sequence.rs

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

competitive/data_structure/splay_tree/
sequence.rs
