haar_lib/ds/
lazy_segtree.rs

1//! モノイド列の区間更新・区間取得($O(\log n)$, $O(\log n)$)ができる。
2use crate::algebra::{act::Act, traits::*};
3use crate::misc::range::range_bounds_to_range;
4use std::ops::RangeBounds;
5
6/// モノイド列の区間更新・区間取得($O(\log n)$, $O(\log n)$)ができる。
7pub struct LazySegtree<M: Monoid, A: Act<M>> {
8    monoid: M,
9    act: A,
10    size: usize,
11    original_size: usize,
12    data: Vec<M::Element>,
13    lazy: Vec<A::Element>,
14}
15
16impl<M, A> LazySegtree<M, A>
17where
18    M: Monoid,
19    A: Act<M>,
20    M::Element: Clone,
21    A::Element: Clone + PartialEq,
22{
23    /// 長さ`n`の[`LazySegtree`]を生成する。
24    pub fn new(monoid: M, act: A, n: usize) -> Self {
25        let size = n.next_power_of_two() * 2;
26        Self {
27            size,
28            original_size: n,
29            data: vec![monoid.id(); size],
30            lazy: vec![act.monoid().id(); size],
31            monoid,
32            act,
33        }
34    }
35
36    /// [`Vec`]から[`LazySegtree`]を構築する。
37    ///
38    /// **Time complexity** $O(|s|)$
39    pub fn from_vec(monoid: M, act: A, s: Vec<M::Element>) -> Self {
40        let n = s.len();
41        let size = n.next_power_of_two() * 2;
42        let mut this = Self {
43            size,
44            original_size: n,
45            data: vec![monoid.id(); size],
46            lazy: vec![act.id(); size],
47            monoid,
48            act,
49        };
50
51        for (i, x) in s.into_iter().enumerate() {
52            this.data[size / 2 + i] = x;
53        }
54
55        for i in (1..size / 2).rev() {
56            this.data[i] = this
57                .monoid
58                .op(this.data[i << 1].clone(), this.data[(i << 1) | 1].clone());
59        }
60
61        this
62    }
63
64    /// 遅延操作を完了させたモノイド列をスライスで返す。
65    ///
66    /// **Time complexity** $O(n)$
67    pub fn to_slice(&mut self) -> &[M::Element] {
68        for i in 1..self.size {
69            self.propagate(i);
70        }
71
72        &self.data[self.size / 2..self.size / 2 + self.original_size]
73    }
74
75    fn propagate(&mut self, i: usize) {
76        if self.lazy[i] == self.act.id() {
77            return;
78        }
79        if i < self.size / 2 {
80            let l = i << 1;
81            let r = (i << 1) | 1;
82
83            self.lazy[l] = self.act.op(self.lazy[l].clone(), self.lazy[i].clone());
84            self.lazy[r] = self.act.op(self.lazy[r].clone(), self.lazy[i].clone());
85        }
86        let len = (self.size / 2) >> (31 - (i as u32).leading_zeros());
87        self.data[i] = self.act.act(
88            &self.monoid,
89            self.data[i].clone(),
90            self.lazy[i].clone(),
91            len,
92        );
93        self.lazy[i] = self.act.id();
94    }
95
96    fn propagate_top_down(&mut self, mut i: usize) {
97        let mut temp = vec![i];
98        while i > 1 {
99            i >>= 1;
100            temp.push(i);
101        }
102
103        for i in temp.into_iter().rev() {
104            self.propagate(i);
105        }
106    }
107
108    fn bottom_up(&mut self, mut i: usize) {
109        while i > 1 {
110            i >>= 1;
111            self.propagate(i << 1);
112            self.propagate((i << 1) | 1);
113            self.data[i] = self
114                .monoid
115                .op(self.data[i << 1].clone(), self.data[(i << 1) | 1].clone());
116        }
117    }
118
119    /// `i`番目の値を返す。
120    pub fn get(&mut self, i: usize) -> M::Element {
121        self.propagate_top_down(i + self.size / 2);
122        self.data[i + self.size / 2].clone()
123    }
124
125    /// 区間`range`で計算を集約して返す。
126    pub fn fold(&mut self, range: impl RangeBounds<usize>) -> M::Element {
127        let (l, r) = range_bounds_to_range(range, 0, self.original_size);
128
129        self.propagate_top_down(l + self.size / 2);
130        if r < self.size / 2 {
131            self.propagate_top_down(r + self.size / 2);
132        }
133
134        let mut ret_l = self.monoid.id();
135        let mut ret_r = self.monoid.id();
136
137        let mut l = l + self.size / 2;
138        let mut r = r + self.size / 2;
139
140        while l < r {
141            if r & 1 == 1 {
142                r -= 1;
143                self.propagate(r);
144                ret_r = self.monoid.op(self.data[r].clone(), ret_r.clone());
145            }
146            if l & 1 == 1 {
147                self.propagate(l);
148                ret_l = self.monoid.op(ret_l.clone(), self.data[l].clone());
149                l += 1;
150            }
151            r >>= 1;
152            l >>= 1;
153        }
154
155        self.monoid.op(ret_l, ret_r)
156    }
157
158    /// `i`番目の値を`value`で置き換える。
159    pub fn assign(&mut self, i: usize, value: M::Element) {
160        self.propagate_top_down(i + self.size / 2);
161        self.data[i + self.size / 2] = value;
162        self.bottom_up(i + self.size / 2);
163    }
164
165    /// 区間`range`を値`x`で更新する。
166    pub fn update(&mut self, range: impl RangeBounds<usize>, x: A::Element) {
167        let (l, r) = range_bounds_to_range(range, 0, self.original_size);
168
169        self.propagate_top_down(l + self.size / 2);
170        if r < self.size / 2 {
171            self.propagate_top_down(r + self.size / 2);
172        }
173
174        {
175            let mut l = l + self.size / 2;
176            let mut r = r + self.size / 2;
177
178            while l < r {
179                if r & 1 == 1 {
180                    r -= 1;
181                    self.lazy[r] = self.act.op(self.lazy[r].clone(), x.clone());
182                }
183                if l & 1 == 1 {
184                    self.lazy[l] = self.act.op(self.lazy[l].clone(), x.clone());
185                    l += 1;
186                }
187                r >>= 1;
188                l >>= 1;
189            }
190        }
191
192        self.bottom_up(l + self.size / 2);
193        if r < self.size / 2 {
194            self.bottom_up(r + self.size / 2);
195        }
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202    use my_testtools::*;
203    use rand::Rng;
204
205    fn test<M, A, F>(n: usize, q: usize, monoid: M, act: A, mut gen: F)
206    where
207        M: Monoid<Element: Copy + PartialEq + std::fmt::Debug> + Clone,
208        A: Act<M, Element: Copy + PartialEq> + Clone,
209        F: FnMut() -> A::Element,
210    {
211        let mut seg = LazySegtree::new(monoid.clone(), act.clone(), n);
212        let mut vec = vec![monoid.id(); n];
213
214        let mut rng = rand::thread_rng();
215
216        for _ in 0..q {
217            let lr = rand_range(&mut rng, 0..n);
218
219            match rng.gen::<u32>() % 2 {
220                0 => {
221                    let x = gen();
222
223                    seg.update(lr.clone(), x);
224                    vec[lr]
225                        .iter_mut()
226                        .for_each(|y| *y = act.act_one(&monoid, *y, x));
227                }
228                1 => {
229                    assert_eq!(
230                        seg.fold(lr.clone()),
231                        vec[lr].iter().cloned().fold_m(&monoid)
232                    );
233                }
234                _ => unreachable!(),
235            }
236        }
237    }
238
239    use crate::algebra;
240
241    #[test]
242    fn add_sum() {
243        let mut rng = rand::thread_rng();
244        let m = algebra::sum::Sum::<u64>::new();
245        test(100, 100, m, algebra::act::add_sum::AddSum(m), || {
246            rng.gen_range(0..1000)
247        });
248    }
249
250    #[test]
251    fn chmax_max() {
252        let mut rng = rand::thread_rng();
253        let m = algebra::min_max::Max::<i64>::new();
254        test(100, 100, m, algebra::act::chmax_max::ChmaxMax(m), || {
255            rng.gen_range(-1000..1000)
256        });
257    }
258
259    #[test]
260    fn chmin_min() {
261        let mut rng = rand::thread_rng();
262        let m = algebra::min_max::Min::<i64>::new();
263        test(100, 100, m, algebra::act::chmin_min::ChminMin(m), || {
264            rng.gen_range(-1000..1000)
265        });
266    }
267}
haar_lib/ds/lazy_segtree.rs

haar_lib/ds/
lazy_segtree.rs
