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

haar_lib/ds/
lazy_splay_tree.rs
