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

haar_lib/ds/
dual_segtree.rs
