haar_lib/ds/
lazy_splay_tree.rs

1//! 遅延スプレー木
2//!
3//! # Problems
4//! - <https://judge.yosupo.jp/problem/dynamic_sequence_range_affine_range_sum>
5
6use std::cell::Cell;
7use std::cmp::Ordering;
8use std::ops::Range;
9use std::ptr;
10
11use crate::algebra::act::*;
12
13#[derive(Clone)]
14struct Manager<M, A> {
15    monoid: M,
16    act: A,
17}
18
19struct Node<M: Monoid, A: Act<M>> {
20    value: M::Element,
21    sum: M::Element,
22    lazy: A::Element,
23    size: usize,
24    rev: bool,
25    lc: *mut Self,
26    rc: *mut Self,
27    par: *mut Self,
28}
29
30impl<M: Monoid, A: Act<M>> Manager<M, A>
31where
32    M::Element: Clone,
33    A::Element: Clone + PartialEq,
34{
35    fn new(monoid: M, act: A) -> Self {
36        Self { monoid, act }
37    }
38
39    fn create(&self, value: M::Element) -> Node<M, A> {
40        Node {
41            value,
42            sum: self.monoid.id(),
43            lazy: self.act.id(),
44            size: 1,
45            rev: false,
46            lc: ptr::null_mut(),
47            rc: ptr::null_mut(),
48            par: ptr::null_mut(),
49        }
50    }
51
52    fn set_value(&self, this: *mut Node<M, A>, value: M::Element) {
53        if !this.is_null() {
54            unsafe {
55                (*this).value = value;
56            }
57        }
58    }
59
60    fn rotate(&self, this: *mut Node<M, A>) {
61        let p = self.par_of(this).unwrap();
62        let pp = self.par_of(p).unwrap();
63
64        self.pushdown(this);
65
66        if self.left_of(p).unwrap() == this {
67            let c = self.right_of(this).unwrap();
68            self.set_left(p, c);
69            self.set_right(this, p);
70        } else {
71            let c = self.left_of(this).unwrap();
72            self.set_right(p, c);
73            self.set_left(this, p);
74        }
75
76        if !pp.is_null() {
77            let pp = unsafe { &mut *pp };
78            if pp.lc == p {
79                pp.lc = this;
80            }
81            if pp.rc == p {
82                pp.rc = this;
83            }
84        }
85
86        assert!(!this.is_null());
87        let this = unsafe { &mut *this };
88        this.par = pp;
89
90        assert!(!p.is_null());
91        let p = unsafe { &mut *p };
92
93        std::mem::swap(&mut this.sum, &mut p.sum);
94        std::mem::swap(&mut this.lazy, &mut p.lazy);
95
96        self.update(p);
97    }
98
99    fn status(&self, this: *mut Node<M, A>) -> i32 {
100        let par = self.par_of(this).unwrap();
101        if par.is_null() {
102            return 0;
103        }
104        let par = unsafe { &*par };
105        if par.lc == this {
106            return 1;
107        }
108        if par.rc == this {
109            return -1;
110        }
111
112        unreachable!()
113    }
114
115    fn reverse(&self, this: *mut Node<M, A>) {
116        if !this.is_null() {
117            unsafe {
118                (*this).rev ^= true;
119            }
120        }
121    }
122
123    fn pushdown(&self, this: *mut Node<M, A>) {
124        if !this.is_null() {
125            let this = unsafe { &mut *this };
126
127            if this.rev {
128                std::mem::swap(&mut this.lc, &mut this.rc);
129                self.reverse(this.lc);
130                self.reverse(this.rc);
131                this.rev = false;
132            }
133
134            if !self.act.monoid().is_id(&this.lazy) {
135                if !this.lc.is_null() {
136                    let lc = unsafe { &mut *this.lc };
137                    lc.lazy = self.act.op(lc.lazy.clone(), this.lazy.clone());
138                }
139
140                if !this.rc.is_null() {
141                    let rc = unsafe { &mut *this.rc };
142                    rc.lazy = self.act.op(rc.lazy.clone(), this.lazy.clone());
143                }
144
145                this.value = self
146                    .act
147                    .act(&self.monoid, this.value.clone(), this.lazy.clone(), 1);
148                this.sum =
149                    self.act
150                        .act(&self.monoid, this.sum.clone(), this.lazy.clone(), this.size);
151                this.lazy = self.act.id();
152            }
153        }
154    }
155
156    fn update(&self, this: *mut Node<M, A>) {
157        assert!(!this.is_null());
158
159        let this = unsafe { &mut *this };
160        self.pushdown(this.lc);
161        self.pushdown(this.rc);
162
163        this.size = 1 + self.size_of(this.lc) + self.size_of(this.rc);
164
165        this.sum = this.value.clone();
166
167        if !this.lc.is_null() {
168            let lc = unsafe { &mut *this.lc };
169            this.sum = self.monoid.op(this.sum.clone(), lc.sum.clone());
170        }
171
172        if !this.rc.is_null() {
173            let rc = unsafe { &mut *this.rc };
174            this.sum = self.monoid.op(this.sum.clone(), rc.sum.clone());
175        }
176    }
177
178    fn splay(&self, this: *mut Node<M, A>) {
179        while self.status(this) != 0 {
180            let par = self.par_of(this).unwrap();
181
182            if self.status(par) == 0 {
183                self.rotate(this);
184            } else if self.status(this) == self.status(par) {
185                self.rotate(par);
186                self.rotate(this);
187            } else {
188                self.rotate(this);
189                self.rotate(this);
190            }
191        }
192    }
193
194    fn get(&self, root: *mut Node<M, A>, mut index: usize) -> *mut Node<M, A> {
195        if root.is_null() {
196            return root;
197        }
198
199        let mut cur = root;
200
201        loop {
202            self.pushdown(cur);
203
204            let left = self.left_of(cur).unwrap();
205            let lsize = self.size_of(left);
206
207            match index.cmp(&lsize) {
208                Ordering::Less => {
209                    cur = left;
210                }
211                Ordering::Equal => {
212                    self.splay(cur);
213                    return cur;
214                }
215                Ordering::Greater => {
216                    cur = self.right_of(cur).unwrap();
217                    index -= lsize + 1;
218                }
219            }
220        }
221    }
222
223    fn merge(&self, left: *mut Node<M, A>, right: *mut Node<M, A>) -> *mut Node<M, A> {
224        if left.is_null() {
225            return right;
226        }
227        if right.is_null() {
228            return left;
229        }
230
231        let cur = self.get(left, self.size_of(left) - 1);
232
233        self.set_right(cur, right);
234
235        self.update(right);
236        self.update(cur);
237
238        cur
239    }
240
241    fn split(&self, root: *mut Node<M, A>, index: usize) -> (*mut Node<M, A>, *mut Node<M, A>) {
242        if root.is_null() {
243            return (ptr::null_mut(), ptr::null_mut());
244        }
245        if index >= self.size_of(root) {
246            return (root, ptr::null_mut());
247        }
248
249        let cur = self.get(root, index);
250        let left = self.left_of(cur).unwrap();
251
252        if !left.is_null() {
253            self.set_par(left, ptr::null_mut());
254            self.update(left);
255        }
256        self.set_left(cur, ptr::null_mut());
257        self.update(cur);
258
259        (left, cur)
260    }
261
262    fn set_left(&self, this: *mut Node<M, A>, left: *mut Node<M, A>) {
263        assert!(!this.is_null());
264        unsafe { (*this).lc = left };
265        if !left.is_null() {
266            unsafe { (*left).par = this };
267        }
268    }
269
270    fn set_right(&self, this: *mut Node<M, A>, right: *mut Node<M, A>) {
271        assert!(!this.is_null());
272        unsafe { (*this).rc = right };
273        if !right.is_null() {
274            unsafe { (*right).par = this };
275        }
276    }
277
278    fn set_par(&self, this: *mut Node<M, A>, par: *mut Node<M, A>) {
279        assert!(!this.is_null());
280        unsafe { (*this).par = par };
281    }
282
283    fn size_of(&self, this: *mut Node<M, A>) -> usize {
284        if this.is_null() {
285            0
286        } else {
287            unsafe { (*this).size }
288        }
289    }
290
291    fn left_of(&self, this: *mut Node<M, A>) -> Option<*mut Node<M, A>> {
292        (!this.is_null()).then(|| unsafe { (*this).lc })
293    }
294
295    fn right_of(&self, this: *mut Node<M, A>) -> Option<*mut Node<M, A>> {
296        (!this.is_null()).then(|| unsafe { (*this).rc })
297    }
298
299    fn par_of(&self, this: *mut Node<M, A>) -> Option<*mut Node<M, A>> {
300        (!this.is_null()).then(|| unsafe { (*this).par })
301    }
302}
303
304/// 遅延スプレー木
305pub struct LazySplayTree<M: Monoid, A: Act<M>> {
306    man: Manager<M, A>,
307    root: Cell<*mut Node<M, A>>,
308}
309
310impl<M: Monoid + Clone, A: Act<M> + Clone> LazySplayTree<M, A>
311where
312    M::Element: Clone,
313    A::Element: Clone + PartialEq,
314{
315    /// `LazySplayTree<A>`を生成
316    pub fn new(monoid: M, act: A) -> Self {
317        let man = Manager::new(monoid, act);
318        let root = Cell::new(ptr::null_mut());
319        Self { man, root }
320    }
321
322    /// 値`value`をもつノード一つのみからなる`SplayTree<M>`を生成
323    pub fn singleton(monoid: M, act: A, value: M::Element) -> Self {
324        let man = Manager::new(monoid, act);
325        let root = Cell::new(Box::into_raw(Box::new(man.create(value))));
326        Self { man, root }
327    }
328
329    fn make_singleton(&self, value: M::Element) -> Self {
330        let man = self.man.clone();
331        let root = Cell::new(Box::into_raw(Box::new(man.create(value))));
332        Self { man, root }
333    }
334
335    /// スプレーツリーの要素数を返す
336    pub fn len(&self) -> usize {
337        self.man.size_of(self.root.get())
338    }
339
340    /// スプレーツリーが要素を持たなければ`true`を返す
341    pub fn is_empty(&self) -> bool {
342        self.root.get().is_null()
343    }
344
345    /// `index`番目の要素の参照を返す
346    pub fn get(&self, index: usize) -> Option<&M::Element> {
347        let node = self.man.get(self.root.get(), index);
348        self.root.set(node);
349
350        (!node.is_null()).then(|| unsafe { &(*node).value })
351    }
352
353    /// `index`番目の要素を`value`に変更する
354    pub fn set(&mut self, index: usize, value: M::Element) {
355        let root = self.man.get(self.root.get(), index);
356        self.man.set_value(root, value);
357        self.man.update(root);
358        self.root.set(root);
359    }
360
361    /// 右側にスプレーツリーを結合する
362    pub fn merge_right(&mut self, right: Self) {
363        let root = self.man.merge(self.root.get(), right.root.get());
364        right.root.set(ptr::null_mut());
365        self.root.set(root);
366    }
367
368    /// 左側にスプレーツリーを結合する
369    pub fn merge_left(&mut self, left: Self) {
370        let root = self.man.merge(left.root.get(), self.root.get());
371        left.root.set(ptr::null_mut());
372        self.root.set(root);
373    }
374
375    /// 左側に`index`個の要素があるように、左右で分割する
376    pub fn split(self, index: usize) -> (Self, Self) {
377        let (l, r) = self.man.split(self.root.get(), index);
378        self.root.set(ptr::null_mut());
379        (
380            Self {
381                man: self.man.clone(),
382                root: Cell::new(l),
383            },
384            Self {
385                man: self.man,
386                root: Cell::new(r),
387            },
388        )
389    }
390
391    /// 要素を`index`番目になるように挿入する
392    pub fn insert(&mut self, index: usize, value: M::Element) {
393        let (l, r) = self.man.split(self.root.get(), index);
394        let node = Box::into_raw(Box::new(self.man.create(value)));
395        let root = self.man.merge(l, self.man.merge(node, r));
396        self.root.set(root);
397    }
398
399    /// `index`番目の要素を削除して、値を返す
400    pub fn remove(&mut self, index: usize) -> Option<M::Element> {
401        let (l, r) = self.man.split(self.root.get(), index);
402        let (m, r) = self.man.split(r, 1);
403
404        self.root.set(self.man.merge(l, r));
405
406        (!m.is_null()).then(|| {
407            let m = unsafe { Box::from_raw(m) };
408            m.value
409        })
410    }
411
412    fn range(
413        &self,
414        start: usize,
415        end: usize,
416    ) -> (*mut Node<M, A>, *mut Node<M, A>, *mut Node<M, A>) {
417        let (m, r) = self.man.split(self.root.get(), end);
418        let (l, m) = self.man.split(m, start);
419        (l, m, r)
420    }
421
422    /// `start..end`の範囲を反転させる
423    pub fn reverse(&mut self, Range { start, end }: Range<usize>) {
424        let (l, m, r) = self.range(start, end);
425        self.man.reverse(m);
426        self.root.set(self.man.merge(self.man.merge(l, m), r));
427    }
428
429    /// `start..end`の範囲でのモノイドの演算の結果を返す
430    pub fn fold(&self, Range { start, end }: Range<usize>) -> M::Element {
431        let (l, m, r) = self.range(start, end);
432        let ret = if m.is_null() {
433            self.man.monoid.id()
434        } else {
435            unsafe { (*m).sum.clone() }
436        };
437        self.root.set(self.man.merge(self.man.merge(l, m), r));
438        ret
439    }
440
441    /// `start..end`の範囲にモノイドの演算を施す。
442    pub fn update(&self, Range { start, end }: Range<usize>, lazy: A::Element) {
443        let (l, m, r) = self.range(start, end);
444        if !m.is_null() {
445            unsafe {
446                (*m).lazy = lazy;
447            }
448        };
449        self.root.set(self.man.merge(self.man.merge(l, m), r));
450    }
451
452    /// 先頭に値を追加する
453    pub fn push_first(&mut self, value: M::Element) {
454        let left = self.make_singleton(value);
455        self.merge_left(left);
456    }
457    /// 末尾に値を追加する
458    pub fn push_last(&mut self, value: M::Element) {
459        let right = self.make_singleton(value);
460        self.merge_right(right);
461    }
462    /// 先頭の値を削除する
463    pub fn pop_first(&mut self) -> Option<M::Element> {
464        self.remove(0)
465    }
466    /// 末尾の値を削除する
467    pub fn pop_last(&mut self) -> Option<M::Element> {
468        self.remove(self.len().checked_sub(1)?)
469    }
470}
471
472#[cfg(test)]
473mod tests {
474    use crate::algebra::{act::add_sum::AddSum, sum::Sum};
475
476    use super::*;
477
478    #[test]
479    fn test_empty() {
480        let monoid = Sum::<u64>::new();
481        let act = AddSum(monoid);
482        let mut s = LazySplayTree::new(monoid, act);
483
484        assert!(s.pop_first().is_none());
485        assert!(s.pop_last().is_none());
486    }
487}
haar_lib/ds/lazy_splay_tree.rs

haar_lib/ds/
lazy_splay_tree.rs
