haar_lib/ds/
dual_segtree.rs

1//! モノイド列の区間更新・点取得($O(\log n)$, $O(\log n)$)ができる。
2#![allow(clippy::wrong_self_convention)]
3
4pub use crate::algebra::traits::Monoid;
5use crate::misc::range::range_bounds_to_range;
6use std::cell::RefCell;
7use std::ops::RangeBounds;
8
9/// モノイド列の区間更新・点取得($O(\log n)$, $O(\log n)$)ができる。
10#[derive(Clone)]
11pub struct DualSegtree<M: Monoid> {
12    monoid: M,
13    original_size: usize,
14    size: usize,
15    data: RefCell<Vec<M::Element>>,
16}
17
18impl<M: Monoid> DualSegtree<M>
19where
20    M::Element: Clone,
21{
22    /// **Time complexity** $O(n)$
23    pub fn new(monoid: M, n: usize) -> Self {
24        let size = n.next_power_of_two() * 2;
25        let data = RefCell::new(vec![monoid.id(); size]);
26        Self {
27            monoid,
28            original_size: n,
29            size,
30            data,
31        }
32    }
33
34    /// モノイド列から`DualSegtree`を構築する。
35    ///
36    /// **Time complexity** $O(|a|)$
37    pub fn from_vec(monoid: M, a: Vec<M::Element>) -> Self {
38        let size = a.len().next_power_of_two() * 2;
39        let original_size = a.len();
40        let mut data = vec![monoid.id(); size];
41        for (i, e) in a.into_iter().enumerate() {
42            data[i + size / 2] = e.clone();
43        }
44        Self {
45            monoid,
46            original_size,
47            size,
48            data: RefCell::new(data),
49        }
50    }
51
52    fn propagate(&self, i: usize) {
53        let mut data = self.data.borrow_mut();
54
55        if i < self.size / 2 {
56            data[i << 1] = self.monoid.op(data[i << 1].clone(), data[i].clone());
57            data[(i << 1) | 1] = self.monoid.op(data[(i << 1) | 1].clone(), data[i].clone());
58            data[i] = self.monoid.id();
59        }
60    }
61
62    fn propagate_top_down(&self, mut i: usize) {
63        let mut temp = vec![];
64        while i > 1 {
65            i >>= 1;
66            temp.push(i);
67        }
68
69        for &i in temp.iter().rev() {
70            self.propagate(i);
71        }
72    }
73
74    /// **Time complexity** $O(\log n)$
75    pub fn get(&self, i: usize) -> M::Element {
76        self.propagate_top_down(i + self.size / 2);
77        self.data.borrow()[i + self.size / 2].clone()
78    }
79
80    /// 遅延操作を完了させたモノイド列を`Vec`で返す。
81    pub fn to_vec(&self) -> Vec<M::Element> {
82        for i in 1..self.size {
83            self.propagate(i);
84        }
85
86        self.data.borrow()[self.size / 2..self.size / 2 + self.original_size].to_vec()
87    }
88
89    /// **Time complexity** $O(\log n)$
90    pub fn update(&mut self, range: impl RangeBounds<usize>, value: M::Element) {
91        let (l, r) = range_bounds_to_range(range, 0, self.original_size);
92
93        let mut l = l + self.size / 2;
94        let mut r = r + self.size / 2;
95
96        self.propagate_top_down(l);
97        self.propagate_top_down(r);
98
99        let mut data = self.data.borrow_mut();
100
101        while l < r {
102            if (r & 1) == 1 {
103                r -= 1;
104                data[r] = self.monoid.op(data[r].clone(), value.clone());
105            }
106            if (l & 1) == 1 {
107                data[l] = self.monoid.op(data[l].clone(), value.clone());
108                l += 1;
109            }
110            l >>= 1;
111            r >>= 1;
112        }
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119    use crate::algebra::sum::*;
120    use my_testtools::*;
121    use rand::Rng;
122
123    #[test]
124    fn test() {
125        let mut rng = rand::thread_rng();
126        let n = 100;
127
128        let m = Sum::<u32>::new();
129
130        let mut a = std::iter::repeat_with(|| rng.gen_range(0..10000))
131            .take(n)
132            .collect::<Vec<_>>();
133        let mut seg = DualSegtree::from_vec(m, a.clone());
134
135        for _ in 0..100 {
136            let lr = rand_range(&mut rng, 0..n);
137            let x = rng.gen_range(0..10000);
138
139            seg.update(lr.clone(), x);
140            a[lr].iter_mut().for_each(|e| m.op_assign_r(e, x));
141
142            assert_eq!(a, seg.to_vec());
143        }
144    }
145}
haar_lib/ds/dual_segtree.rs

haar_lib/ds/
dual_segtree.rs
