haar_lib/ds/
dynamic_lazy_segtree.rs

1//! 動的遅延セグメント木
2use crate::algebra::{act::Act, traits::*};
3use std::ops::Range;
4use std::ptr;
5
6#[derive(Clone, Debug)]
7struct Node<M: Monoid, A: Act<M>> {
8    value: M::Element,
9    lazy: A::Element,
10    left: *mut Self,
11    right: *mut Self,
12}
13
14impl<M: Monoid, A: Act<M>> Node<M, A> {
15    fn new(monoid: &M, act: &A) -> Self {
16        Self {
17            value: monoid.id(),
18            lazy: act.id(),
19            left: ptr::null_mut(),
20            right: ptr::null_mut(),
21        }
22    }
23}
24
25/// 動的遅延セグメント木
26#[derive(Clone, Debug)]
27pub struct DynamicLazySegtree<M: Monoid, A: Act<M>> {
28    monoid: M,
29    act: A,
30    root: *mut Node<M, A>,
31    to: usize,
32}
33
34impl<M: Monoid, A: Act<M>> DynamicLazySegtree<M, A> {
35    /// `DynamicLazySegtree<A>`を生成する。
36    pub fn new(monoid: M, act: A) -> Self {
37        Self {
38            root: Box::into_raw(Box::new(Node::new(&monoid, &act))),
39            to: 1,
40            monoid,
41            act,
42        }
43    }
44}
45
46impl<M: Monoid, A: Act<M>> DynamicLazySegtree<M, A>
47where
48    M::Element: Clone,
49    A::Element: Clone + PartialEq,
50{
51    fn _propagate(&self, t: *mut Node<M, A>, from: usize, to: usize) {
52        assert!(!t.is_null());
53        let t = unsafe { &mut *t };
54
55        if t.lazy == self.act.id() {
56            return;
57        }
58        if to - from > 1 {
59            if t.left.is_null() {
60                t.left = Box::into_raw(Box::new(Node::new(&self.monoid, &self.act)));
61            }
62            let left = unsafe { &mut *t.left };
63            left.lazy = self.act.op(left.lazy.clone(), t.lazy.clone());
64
65            if t.right.is_null() {
66                t.right = Box::into_raw(Box::new(Node::new(&self.monoid, &self.act)));
67            }
68            let right = unsafe { &mut *t.right };
69            right.lazy = self.act.op(right.lazy.clone(), t.lazy.clone());
70        }
71        let len = to - from;
72        t.value = self
73            .act
74            .act(&self.monoid, t.value.clone(), t.lazy.clone(), len);
75        t.lazy = self.act.id();
76    }
77
78    fn _update(
79        &self,
80        mut cur: *mut Node<M, A>,
81        from: usize,
82        to: usize,
83        s: usize,
84        t: usize,
85        value: A::Element,
86    ) -> *mut Node<M, A> {
87        if cur.is_null() {
88            cur = Box::into_raw(Box::new(Node::new(&self.monoid, &self.act)));
89        }
90        let cur = unsafe { &mut *cur };
91
92        self._propagate(cur, from, to);
93
94        if to - from == 1 {
95            if s <= from && to <= t {
96                cur.lazy = self.act.op(cur.lazy.clone(), value);
97            }
98            self._propagate(cur, from, to);
99            return cur;
100        }
101
102        if to < s || t < from {
103            return cur;
104        }
105        if s <= from && to <= t {
106            cur.lazy = self.act.op(cur.lazy.clone(), value);
107            self._propagate(cur, from, to);
108            return cur;
109        }
110
111        let mid = (from + to) / 2;
112        cur.left = self._update(cur.left, from, mid, s, t, value.clone());
113        cur.right = self._update(cur.right, mid, to, s, t, value);
114        assert!(!cur.left.is_null() && !cur.right.is_null());
115        let left = unsafe { &*cur.left };
116        let right = unsafe { &*cur.right };
117        cur.value = self.monoid.op(left.value.clone(), right.value.clone());
118        cur
119    }
120
121    /// 範囲`s..t`を`value`で更新する。
122    pub fn update(&mut self, Range { start: s, end: t }: Range<usize>, value: A::Element) {
123        loop {
124            if t <= self.to {
125                break;
126            }
127            self.to *= 2;
128
129            let mut new_root = Box::new(Node::new(&self.monoid, &self.act));
130            new_root.left = self.root;
131
132            self.root = Box::into_raw(new_root);
133        }
134
135        self._update(self.root, 0, self.to, s, t, value);
136    }
137
138    fn _fold(
139        &self,
140        cur: *mut Node<M, A>,
141        from: usize,
142        to: usize,
143        s: usize,
144        t: usize,
145    ) -> M::Element {
146        if cur.is_null() {
147            return self.monoid.id();
148        }
149
150        let cur = unsafe { &mut *cur };
151
152        self._propagate(cur, from, to);
153        if to <= s || t <= from {
154            return self.monoid.id();
155        }
156        if s <= from && to <= t {
157            return cur.value.clone();
158        }
159
160        let mid = (from + to) / 2;
161        let lv = self._fold(cur.left, from, mid, s, t);
162        let rv = self._fold(cur.right, mid, to, s, t);
163
164        self.monoid.op(lv, rv)
165    }
166
167    /// 範囲`s..t`で計算を集約する。
168    pub fn fold(&mut self, Range { start: s, end: t }: Range<usize>) -> M::Element {
169        self._fold(self.root, 0, self.to, s, t)
170    }
171}
haar_lib/ds/dynamic_lazy_segtree.rs

haar_lib/ds/
dynamic_lazy_segtree.rs
