cp_library_rs/data_structure/
lazy_segment_tree.rs

1//! 遅延評価セグメント木
2//! - 参考: <https://drken1215.hatenablog.com/entry/2024/11/17/035045>
3
4use crate::{
5    algebraic_structure::actedmonoid::ActedMonoid, tree::show_binary_tree::ShowBinaryTree,
6};
7use core::fmt;
8use std::{
9    fmt::Debug,
10    ops::{
11        Bound::{Excluded, Included, Unbounded},
12        RangeBounds,
13    },
14};
15
16/// 遅延評価セグメント木
17#[derive(Debug)]
18pub struct LazySegmentTree<M: ActedMonoid> {
19    pub size: usize,
20    log: usize,
21    offset: usize,
22    data: Vec<M::Val>,
23    lazy: Vec<M::Act>,
24}
25
26impl<M: ActedMonoid> LazySegmentTree<M> {
27    #[inline]
28    fn parse_range<R: RangeBounds<usize>>(&self, range: &R) -> Option<(usize, usize)> {
29        let start = match range.start_bound() {
30            Unbounded => 0,
31            Excluded(&v) => v + 1,
32            Included(&v) => v,
33        };
34        let end = match range.end_bound() {
35            Unbounded => self.size,
36            Excluded(&v) => v,
37            Included(&v) => v + 1,
38        };
39        if start <= end && end <= self.size {
40            Some((start, end))
41        } else {
42            None
43        }
44    }
45
46    /// 遅延評価セグメント木を初期化する
47    /// - `n`: 配列サイズ
48    pub fn new(n: usize) -> Self {
49        assert!(n > 0);
50        let offset = n.next_power_of_two();
51        let log = offset.trailing_zeros() as usize;
52        Self {
53            size: n,
54            log,
55            offset,
56            data: vec![M::e(); offset << 1],
57            lazy: vec![M::id(); offset << 1],
58        }
59    }
60
61    /// 配列から構築する
62    pub fn from_vec(src: Vec<M::Val>) -> Self {
63        let mut seg = Self::new(src.len());
64        for (i, v) in src.into_iter().enumerate() {
65            seg.data[seg.offset + i] = v;
66        }
67        for k in (1..seg.offset).rev() {
68            seg.pull_dat(k);
69        }
70        seg
71    }
72
73    /// 1つ上の情報を更新
74    fn pull_dat(&mut self, k: usize) {
75        self.data[k] = M::op(&self.data[k << 1], &self.data[k << 1 | 1]);
76    }
77
78    /// 作用をノードに反映
79    fn apply_lazy(&mut self, k: usize, f: &M::Act) {
80        self.data[k] = M::mapping(&self.data[k], f);
81        if k < self.offset {
82            self.lazy[k] = M::compose(&self.lazy[k], f);
83        }
84    }
85
86    /// 遅延値を子へ伝播
87    fn push_lazy(&mut self, k: usize) {
88        if self.lazy[k] == M::id() {
89            return;
90        }
91        let f = self.lazy[k].clone();
92        self.apply_lazy(k << 1, &f);
93        self.apply_lazy(k << 1 | 1, &f);
94        self.lazy[k] = M::id();
95    }
96
97    /// 葉から親方向へ更新
98    fn pull_dat_deep(&mut self, k: usize) {
99        for h in 1..=self.log {
100            self.pull_dat(k >> h);
101        }
102    }
103
104    /// 根から葉方向へ遅延評価
105    fn push_lazy_deep(&mut self, k: usize) {
106        for h in (1..=self.log).rev() {
107            self.push_lazy(k >> h);
108        }
109    }
110
111    /// 要素を更新する
112    pub fn set(&mut self, i: usize, v: M::Val) {
113        assert!(i < self.size);
114        let k = i + self.offset;
115        self.push_lazy_deep(k);
116        self.data[k] = v;
117        self.pull_dat_deep(k);
118    }
119
120    /// 要素を取得する
121    pub fn get_at(&mut self, i: usize) -> M::Val {
122        assert!(i < self.size);
123        let k = i + self.offset;
124        self.push_lazy_deep(k);
125        self.data[k].clone()
126    }
127
128    /// 要素へ作用させる
129    pub fn apply_at(&mut self, i: usize, f: M::Act) {
130        assert!(i < self.size);
131        let k = i + self.offset;
132        self.push_lazy_deep(k);
133        self.data[k] = M::mapping(&self.data[k], &f);
134        self.pull_dat_deep(k);
135    }
136
137    /// 区間に`val`を作用させる
138    /// - `range`: `[left, right)`
139    pub fn apply<R: RangeBounds<usize> + fmt::Debug>(&mut self, range: R, val: M::Act) {
140        let Some((left, right)) = self.parse_range(&range) else {
141            panic!("The given range is wrong: {:?}", range);
142        };
143        self.apply_range(left, right, &val);
144    }
145
146    fn apply_range(&mut self, mut l: usize, mut r: usize, f: &M::Act) {
147        if l == r {
148            return;
149        }
150        l += self.offset;
151        r += self.offset;
152        for h in (1..=self.log).rev() {
153            if ((l >> h) << h) != l {
154                self.push_lazy(l >> h);
155            }
156            if ((r >> h) << h) != r {
157                self.push_lazy((r - 1) >> h);
158            }
159        }
160        let (l0, r0) = (l, r);
161        while l < r {
162            if (l & 1) == 1 {
163                self.apply_lazy(l, f);
164                l += 1;
165            }
166            if (r & 1) == 1 {
167                r -= 1;
168                self.apply_lazy(r, f);
169            }
170            l >>= 1;
171            r >>= 1;
172        }
173        l = l0;
174        r = r0;
175        for h in 1..=self.log {
176            if ((l >> h) << h) != l {
177                self.pull_dat(l >> h);
178            }
179            if ((r >> h) << h) != r {
180                self.pull_dat((r - 1) >> h);
181            }
182        }
183    }
184
185    /// 区間を取得する
186    /// - `range`: `[left, right)`
187    pub fn get<R: RangeBounds<usize> + fmt::Debug>(&mut self, range: R) -> M::Val {
188        let Some((left, right)) = self.parse_range(&range) else {
189            panic!("The given range is wrong: {:?}", range);
190        };
191        self.prod_range(left, right)
192    }
193
194    fn prod_range(&mut self, mut l: usize, mut r: usize) -> M::Val {
195        if l == r {
196            return M::e();
197        }
198        l += self.offset;
199        r += self.offset;
200        for h in (1..=self.log).rev() {
201            if ((l >> h) << h) != l {
202                self.push_lazy(l >> h);
203            }
204            if ((r >> h) << h) != r {
205                self.push_lazy((r - 1) >> h);
206            }
207        }
208        let (mut val_left, mut val_right) = (M::e(), M::e());
209        while l < r {
210            if (l & 1) == 1 {
211                val_left = M::op(&val_left, &self.data[l]);
212                l += 1;
213            }
214            if (r & 1) == 1 {
215                r -= 1;
216                val_right = M::op(&self.data[r], &val_right);
217            }
218            l >>= 1;
219            r >>= 1;
220        }
221        M::op(&val_left, &val_right)
222    }
223
224    /// 全区間の集約
225    pub fn all_prod(&self) -> M::Val {
226        self.data[1].clone()
227    }
228
229    /// 左端を固定した2分探索
230    /// - 返り値: (prod([l, Val)), Val)
231    pub fn max_right<Act>(&mut self, l: usize, f: Act) -> (M::Val, usize)
232    where
233        Act: Fn(M::Val) -> bool,
234    {
235        assert!(f(M::e()));
236        if l == self.size {
237            return (M::e(), self.size);
238        }
239        let mut l = l + self.offset;
240        self.push_lazy_deep(l);
241        let mut sum = M::e();
242        loop {
243            while l % 2 == 0 {
244                l >>= 1;
245            }
246            let tmp = M::op(&sum, &self.data[l]);
247            if !f(tmp.clone()) {
248                while l < self.offset {
249                    self.push_lazy(l);
250                    l <<= 1;
251                    let tmp = M::op(&sum, &self.data[l]);
252                    if f(tmp.clone()) {
253                        sum = tmp;
254                        l += 1;
255                    }
256                }
257                return (sum, l - self.offset);
258            }
259            sum = tmp;
260            l += 1;
261            if (l & l.wrapping_neg()) == l {
262                break;
263            }
264        }
265        (sum, self.size)
266    }
267
268    /// 右端を固定した2分探索
269    /// - 返り値: (prod([Val, r)), Val)
270    pub fn min_left<Act>(&mut self, r: usize, f: Act) -> (M::Val, usize)
271    where
272        Act: Fn(M::Val) -> bool,
273    {
274        assert!(f(M::e()));
275        if r == 0 {
276            return (M::e(), 0);
277        }
278        let mut r = r + self.offset;
279        self.push_lazy_deep(r - 1);
280        let mut sum = M::e();
281        loop {
282            r -= 1;
283            while r > 1 && (r % 2) == 1 {
284                r >>= 1;
285            }
286            let tmp = M::op(&self.data[r], &sum);
287            if !f(tmp.clone()) {
288                while r < self.offset {
289                    self.push_lazy(r);
290                    r = r * 2 + 1;
291                    let tmp = M::op(&self.data[r], &sum);
292                    if f(tmp.clone()) {
293                        sum = tmp;
294                        r -= 1;
295                    }
296                }
297                return (sum, r + 1 - self.offset);
298            }
299            sum = tmp;
300            if (r & r.wrapping_neg()) == r {
301                break;
302            }
303        }
304        (sum, 0)
305    }
306}
307
308impl<M: ActedMonoid> FromIterator<M::Val> for LazySegmentTree<M> {
309    fn from_iter<T: IntoIterator<Item = M::Val>>(iter: T) -> Self {
310        // 配列にする
311        let arr: Vec<M::Val> = iter.into_iter().collect();
312        Self::from_vec(arr)
313    }
314}
315
316impl<M> ShowBinaryTree<usize> for LazySegmentTree<M>
317where
318    M: ActedMonoid,
319    M::Act: Debug,
320    M::Val: Debug,
321{
322    fn get_root(&self) -> Option<usize> {
323        Some(1)
324    }
325    fn get_left(&self, &i: &usize) -> Option<usize> {
326        (i * 2 < self.offset * 2).then_some(i * 2)
327    }
328    fn get_right(&self, &i: &usize) -> Option<usize> {
329        (i * 2 + 1 < self.offset * 2).then_some(i * 2 + 1)
330    }
331    fn print_node(&self, &i: &usize) -> String {
332        format!("[data:{:?}, lazy:{:?}]", self.data[i], self.lazy[i])
333    }
334}
cp_library_rs/data_structure/lazy_segment_tree.rs

cp_library_rs/data_structure/
lazy_segment_tree.rs
