competitive/data_structure/
implicit_splay_tree.rs

1use super::{
2    Allocator, LazyMapMonoid, MemoryPool,
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::Split3,
9    },
10    splay_operations,
11};
12use std::{
13    cmp::Ordering,
14    fmt::{self, Debug},
15    marker::PhantomData,
16    mem::{ManuallyDrop, replace},
17    ops::{DerefMut, RangeBounds},
18};
19
20type ImplicitSplayTreeRoot<T> = BstRoot<ImplicitSplayTreeSpec<T>>;
21type ImplicitSplayTreeNode<T> = BstNode<ImplicitSplayTreeData<T>>;
22
23pub struct ImplicitSplayTreeSpec<T> {
24    _marker: PhantomData<fn() -> T>,
25}
26
27pub struct ImplicitSplayTreeData<T>
28where
29    T: LazyMapMonoid,
30{
31    value: LazyMapElement<T>,
32    size: usize,
33    rev: bool,
34}
35
36impl<T> Debug for ImplicitSplayTreeData<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("ImplicitSplayTreeData")
42            .field("value", &self.value)
43            .field("size", &self.size)
44            .field("rev", &self.rev)
45            .finish()
46    }
47}
48
49impl<T> BstDataAccess<data::marker::Size> for ImplicitSplayTreeData<T>
50where
51    T: LazyMapMonoid,
52{
53    type Value = usize;
54
55    fn bst_data(&self) -> &Self::Value {
56        &self.size
57    }
58
59    fn bst_data_mut(&mut self) -> &mut Self::Value {
60        &mut self.size
61    }
62}
63
64impl<T> BstDataAccess<data::marker::LazyMap> for ImplicitSplayTreeData<T>
65where
66    T: LazyMapMonoid,
67{
68    type Value = LazyMapElement<T>;
69
70    fn bst_data(&self) -> &Self::Value {
71        &self.value
72    }
73
74    fn bst_data_mut(&mut self) -> &mut Self::Value {
75        &mut self.value
76    }
77}
78
79impl<T> ImplicitSplayTreeSpec<T>
80where
81    T: LazyMapMonoid,
82{
83    fn update_act(mut node: BstDataMutRef<'_, Self>, act: &T::Act) {
84        if T::is_act_unit(act) {
85            return;
86        }
87        T::act_operate_assign(&mut node.data_mut().value.act, act);
88        node.data_mut().value.key = T::act_key(&node.reborrow().into_data().value.key, act);
89        if let Some(agg) = T::act_agg(&node.reborrow().into_data().value.agg, act) {
90            node.data_mut().value.agg = agg;
91        } else {
92            Self::top_down(node.reborrow_datamut());
93            Self::bottom_up(node);
94        }
95    }
96
97    fn reverse(mut node: BstDataMutRef<'_, Self>) {
98        node.swap_children();
99        let data = node.data_mut();
100        T::toggle(&mut data.value.agg);
101        data.rev ^= true;
102    }
103}
104
105impl<T> BstSpec for ImplicitSplayTreeSpec<T>
106where
107    T: LazyMapMonoid,
108{
109    type Parent = WithNoParent<Self::Data>;
110    type Data = ImplicitSplayTreeData<T>;
111
112    fn top_down(mut node: BstDataMutRef<'_, Self>) {
113        if !T::is_act_unit(&node.reborrow().into_data().value.act) {
114            let act = replace(&mut node.data_mut().value.act, T::act_unit());
115            if let Ok(left) = node.reborrow_datamut().left().descend() {
116                Self::update_act(left, &act);
117            }
118            if let Ok(right) = node.reborrow_datamut().right().descend() {
119                Self::update_act(right, &act);
120            }
121        }
122        if node.reborrow().into_data().rev {
123            node.data_mut().rev = false;
124            if let Ok(left) = node.reborrow_datamut().left().descend() {
125                Self::reverse(left);
126            }
127            if let Ok(right) = node.reborrow_datamut().right().descend() {
128                Self::reverse(right);
129            }
130        }
131    }
132
133    fn bottom_up(mut node: BstDataMutRef<'_, Self>) {
134        let mut agg = T::single_agg(&node.reborrow().into_data().value.key);
135        let mut size = 1;
136        if let Ok(left) = node.reborrow().left().descend() {
137            let data = left.into_data();
138            agg = T::agg_operate(&data.value.agg, &agg);
139            size += data.size;
140        }
141        if let Ok(right) = node.reborrow().right().descend() {
142            let data = right.into_data();
143            agg = T::agg_operate(&agg, &data.value.agg);
144            size += data.size;
145        }
146        let data = node.data_mut();
147        data.value.agg = agg;
148        data.size = size;
149    }
150
151    fn merge(
152        left: Option<ImplicitSplayTreeRoot<T>>,
153        right: Option<ImplicitSplayTreeRoot<T>>,
154    ) -> Option<ImplicitSplayTreeRoot<T>> {
155        splay_operations::merge(left, right)
156    }
157
158    fn split<Seeker>(
159        node: Option<ImplicitSplayTreeRoot<T>>,
160        seeker: Seeker,
161        equal_side: EqualSide,
162    ) -> (
163        Option<ImplicitSplayTreeRoot<T>>,
164        Option<ImplicitSplayTreeRoot<T>>,
165    )
166    where
167        Seeker: BstSeeker<Spec = Self>,
168    {
169        splay_operations::split(node, seeker, equal_side)
170    }
171}
172
173pub struct ImplicitSplayTree<T, A = MemoryPool<ImplicitSplayTreeNode<T>>>
174where
175    T: LazyMapMonoid,
176    A: Allocator<ImplicitSplayTreeNode<T>>,
177{
178    root: Option<ImplicitSplayTreeRoot<T>>,
179    length: usize,
180    allocator: ManuallyDrop<A>,
181    _marker: PhantomData<fn() -> T>,
182}
183
184impl<T, A> Default for ImplicitSplayTree<T, A>
185where
186    T: LazyMapMonoid,
187    A: Allocator<ImplicitSplayTreeNode<T>> + Default,
188{
189    fn default() -> Self {
190        Self {
191            root: None,
192            length: 0,
193            allocator: ManuallyDrop::new(A::default()),
194            _marker: PhantomData,
195        }
196    }
197}
198
199impl<T, A> Drop for ImplicitSplayTree<T, A>
200where
201    T: LazyMapMonoid,
202    A: Allocator<ImplicitSplayTreeNode<T>>,
203{
204    fn drop(&mut self) {
205        unsafe {
206            if let Some(root) = self.root.take() {
207                root.into_dying().drop_all(self.allocator.deref_mut());
208            }
209            ManuallyDrop::drop(&mut self.allocator);
210        }
211    }
212}
213
214impl<T> ImplicitSplayTree<T>
215where
216    T: LazyMapMonoid,
217{
218    pub fn new() -> Self {
219        Self::default()
220    }
221
222    pub fn with_capacity(capacity: usize) -> Self {
223        Self {
224            root: None,
225            length: 0,
226            allocator: ManuallyDrop::new(MemoryPool::with_capacity(capacity)),
227            _marker: PhantomData,
228        }
229    }
230}
231
232impl<T, A> ImplicitSplayTree<T, A>
233where
234    T: LazyMapMonoid,
235    A: Allocator<ImplicitSplayTreeNode<T>>,
236{
237    fn node(&mut self, key: T::Key) -> ImplicitSplayTreeRoot<T> {
238        BstRoot::from_data(
239            ImplicitSplayTreeData {
240                value: LazyMapElement::from_key(key),
241                size: 1,
242                rev: false,
243            },
244            self.allocator.deref_mut(),
245        )
246    }
247
248    #[inline]
249    fn splay<Seeker>(&mut self, seeker: Seeker) -> Option<Ordering>
250    where
251        Seeker: BstSeeker<Spec = ImplicitSplayTreeSpec<T>>,
252    {
253        let (ordering, root) = splay_operations::splay(self.root.take()?, seeker);
254        self.root = Some(root);
255        Some(ordering)
256    }
257
258    pub fn len(&self) -> usize {
259        self.length
260    }
261
262    pub fn is_empty(&self) -> bool {
263        self.length == 0
264    }
265
266    pub fn update<R>(&mut self, range: R, act: T::Act)
267    where
268        R: RangeBounds<usize>,
269    {
270        let mut split = Split3::seek_by_size(&mut self.root, range);
271        if let Some(root) = split.mid_datamut() {
272            ImplicitSplayTreeSpec::update_act(root, &act);
273        }
274    }
275
276    pub fn fold<R>(&mut self, range: R) -> T::Agg
277    where
278        R: RangeBounds<usize>,
279    {
280        let split = Split3::seek_by_size(&mut self.root, range);
281        split
282            .mid()
283            .map(|node| node.into_data().value.agg.clone())
284            .unwrap_or_else(T::agg_unit)
285    }
286
287    pub fn reverse<R>(&mut self, range: R)
288    where
289        R: RangeBounds<usize>,
290    {
291        let mut split = Split3::seek_by_size(&mut self.root, range);
292        if let Some(root) = split.mid_datamut() {
293            ImplicitSplayTreeSpec::reverse(root);
294        }
295    }
296
297    pub fn get(&mut self, index: usize) -> Option<&T::Key> {
298        if index >= self.length {
299            return None;
300        }
301        self.splay(SeekBySize::new(index));
302        Some(&self.root.as_ref()?.reborrow().into_data().value.key)
303    }
304
305    pub fn modify<F>(&mut self, index: usize, f: F)
306    where
307        F: FnOnce(&T::Key) -> T::Key,
308    {
309        assert!(index < self.length);
310        self.splay(SeekBySize::new(index));
311        let mut root = self.root.as_mut().unwrap().borrow_datamut();
312        ImplicitSplayTreeSpec::top_down(root.reborrow_datamut());
313        {
314            let data = root.data_mut();
315            data.value.key = f(&data.value.key);
316        }
317        ImplicitSplayTreeSpec::bottom_up(root);
318    }
319
320    pub fn insert(&mut self, index: usize, key: T::Key) {
321        assert!(index <= self.length);
322        let mut node = self.node(key);
323        if self.root.is_none() {
324            self.root = Some(node);
325        } else if index == self.length {
326            self.splay(SeekBySize::new(index));
327            unsafe { node.borrow_mut().left_mut().set(self.root.take().unwrap()) };
328            ImplicitSplayTreeSpec::bottom_up(node.borrow_datamut());
329            self.root = Some(node);
330        } else {
331            self.splay(SeekBySize::new(index));
332            let mut root = self.root.take().unwrap();
333            let left = unsafe { root.borrow_mut().left_mut().take() };
334            if let Some(left) = left {
335                unsafe { node.borrow_mut().left_mut().set(left) };
336            }
337            ImplicitSplayTreeSpec::bottom_up(root.borrow_datamut());
338            unsafe { node.borrow_mut().right_mut().set(root) };
339            ImplicitSplayTreeSpec::bottom_up(node.borrow_datamut());
340            self.root = Some(node);
341        }
342        self.length += 1;
343    }
344
345    pub fn remove(&mut self, index: usize) -> Option<T::Key> {
346        if index >= self.length {
347            return None;
348        }
349        self.splay(SeekBySize::new(index));
350        let mut node = self.root.take().unwrap();
351        ImplicitSplayTreeSpec::top_down(node.borrow_datamut());
352        let left = unsafe { node.borrow_mut().left_mut().take() };
353        let right = unsafe { node.borrow_mut().right_mut().take() };
354        self.root = ImplicitSplayTreeSpec::merge(left, right);
355        self.length -= 1;
356        let data = unsafe { node.into_dying().into_data(self.allocator.deref_mut()) };
357        Some(data.value.key)
358    }
359
360    pub fn partition_point_acc<F>(&mut self, left: usize, mut pred: F) -> usize
361    where
362        F: FnMut(&T::Agg) -> bool,
363    {
364        let mut split3 = Split3::seek_by_size(&mut self.root, left..);
365        let front_size = split3
366            .left()
367            .map(|node| node.into_data().size)
368            .unwrap_or_default();
369        let split = split3.split_mid(SeekByAccCond::new(|acc| !pred(acc)), EqualSide::Right);
370        let index = split
371            .left()
372            .map(|node| node.into_data().size)
373            .unwrap_or_default();
374        front_size + index
375    }
376
377    pub fn rpartition_point_acc<F>(&mut self, right: usize, mut pred: F) -> usize
378    where
379        F: FnMut(&T::Agg) -> bool,
380    {
381        let mut split3 = Split3::seek_by_size(&mut self.root, ..right);
382        let split = split3.split_mid(SeekByRaccCond::new(|acc| !pred(acc)), EqualSide::Left);
383        split
384            .left()
385            .map(|node| node.into_data().size)
386            .unwrap_or_default()
387    }
388
389    pub fn rotate_left(&mut self, mid: usize) {
390        assert!(mid <= self.length);
391        if mid == 0 || mid == self.length {
392            return;
393        }
394        let (left, right) =
395            ImplicitSplayTreeSpec::split(self.root.take(), SeekBySize::new(mid), EqualSide::Right);
396        self.root = ImplicitSplayTreeSpec::merge(right, left);
397    }
398
399    pub fn rotate_right(&mut self, k: usize) {
400        assert!(k <= self.length);
401        self.rotate_left(self.length - k);
402    }
403}
404
405impl<T, A> Extend<T::Key> for ImplicitSplayTree<T, A>
406where
407    T: LazyMapMonoid,
408    A: Allocator<ImplicitSplayTreeNode<T>>,
409{
410    fn extend<I>(&mut self, iter: I)
411    where
412        I: IntoIterator<Item = T::Key>,
413    {
414        let nodes = iter
415            .into_iter()
416            .map(|key| self.node(key))
417            .collect::<Vec<_>>();
418        let len = nodes.len();
419        let root = if len == 0 {
420            None
421        } else {
422            let mut stack = Vec::with_capacity(64);
423            stack.push((0, len, None::<(usize, usize)>, false));
424            while let Some((start, end, parent, visited)) = stack.pop() {
425                if start == end {
426                    continue;
427                }
428                let mid = start + (end - start) / 2;
429                if visited {
430                    ImplicitSplayTreeSpec::bottom_up(
431                        BstRoot::new(nodes[mid].node).borrow_datamut(),
432                    );
433                    continue;
434                }
435                if let Some((parent, direction)) = parent {
436                    let mut parent = nodes[parent].node;
437                    unsafe { parent.as_mut().child[direction] = Some(nodes[mid].node) };
438                }
439                stack.push((start, end, parent, true));
440                stack.push((mid + 1, end, Some((mid, 1)), false));
441                stack.push((start, mid, Some((mid, 0)), false));
442            }
443            Some(BstRoot::new(nodes[len / 2].node))
444        };
445        self.root = ImplicitSplayTreeSpec::merge(self.root.take(), root);
446        self.length += len;
447    }
448}
449
450#[cfg(test)]
451mod tests {
452    use super::*;
453    use crate::{
454        algebra::RangeMaxRangeUpdate,
455        rand,
456        tools::{NotEmptySegment, Xorshift},
457    };
458
459    #[test]
460    fn test_implicit_splay_tree() {
461        const N: usize = 1_000;
462        const Q: usize = 20_000;
463        const A: i64 = 1_000_000_000;
464
465        let mut rng = Xorshift::default();
466        rand!(rng, mut arr: [-A..A; N]);
467        let mut tree = ImplicitSplayTree::<RangeMaxRangeUpdate<_>>::new();
468        tree.extend(arr.iter().copied());
469        for _ in 0..Q {
470            assert_eq!(arr.len(), tree.len());
471            rand!(rng, ty: 0..10, (l, r): NotEmptySegment(arr.len()));
472            match ty {
473                0 if arr.len() < N * 2 => {
474                    rand!(rng, i: ..=arr.len(), x: -A..A);
475                    tree.insert(i, x);
476                    arr.insert(i, x);
477                }
478                1 if arr.len() > 1 => {
479                    rand!(rng, i: ..arr.len());
480                    assert_eq!(arr.remove(i), tree.remove(i).unwrap());
481                }
482                2 => assert_eq!(tree.fold(l..r), *arr[l..r].iter().max().unwrap()),
483                3 => {
484                    rand!(rng, x: -A..A);
485                    tree.update(l..r, Some(x));
486                    arr[l..r].fill(x);
487                }
488                4 => {
489                    tree.reverse(l..r);
490                    arr[l..r].reverse();
491                }
492                5 => {
493                    rand!(rng, left: ..=arr.len(), x: -A..A);
494                    assert_eq!(
495                        tree.partition_point_acc(left, |&value| value < x),
496                        arr[left..]
497                            .iter()
498                            .scan(i64::MIN, |acc, &value| {
499                                *acc = (*acc).max(value);
500                                Some(*acc)
501                            })
502                            .position(|value| value >= x)
503                            .map_or(arr.len(), |index| left + index),
504                    );
505                }
506                6 => {
507                    rand!(rng, right: ..=arr.len(), x: -A..A);
508                    assert_eq!(
509                        tree.rpartition_point_acc(right, |&value| value < x),
510                        arr[..right]
511                            .iter()
512                            .rev()
513                            .scan(i64::MIN, |acc, &value| {
514                                *acc = (*acc).max(value);
515                                Some(*acc)
516                            })
517                            .position(|value| value >= x)
518                            .map_or(0, |index| right - index),
519                    );
520                }
521                7 => {
522                    rand!(rng, mid: ..=arr.len());
523                    tree.rotate_left(mid);
524                    arr.rotate_left(mid);
525                }
526                8 => {
527                    rand!(rng, count: ..=arr.len());
528                    tree.rotate_right(count);
529                    arr.rotate_right(count);
530                }
531                _ => {
532                    rand!(rng, index: ..arr.len(), value: -A..A);
533                    tree.modify(index, |_| value);
534                    arr[index] = value;
535                }
536            }
537            assert_eq!(tree.get(tree.len()), None);
538            rand!(rng, index: ..arr.len());
539            assert_eq!(tree.get(index), arr.get(index));
540        }
541    }
542}
competitive/data_structure/implicit_splay_tree.rs

competitive/data_structure/
implicit_splay_tree.rs
