haar_lib/ds/
segtree.rs

1//! モノイド列の点更新・区間取得($O(\log n)$, $O(\log n)$)ができる。
2pub use crate::algebra::traits::Monoid;
3use crate::misc::range::range_bounds_to_range;
4use std::ops::{Index, RangeBounds};
5
6/// モノイド列の点更新・区間取得($O(\log n)$, $O(\log n)$)ができる。
7#[derive(Clone)]
8pub struct Segtree<M: Monoid> {
9    original_size: usize,
10    size: usize,
11    data: Vec<M>,
12}
13
14impl<M: Monoid + Clone> Segtree<M> {
15    /// **Time complexity** $O(n)$
16    pub fn new(n: usize) -> Self {
17        let size = n.next_power_of_two() * 2;
18        Segtree {
19            original_size: n,
20            size,
21            data: vec![M::id(); size],
22        }
23    }
24
25    /// モノイド列から`Segtree`を構築する。
26    ///
27    /// **Time complexity** $O(|s|)$
28    pub fn from_vec(s: Vec<M>) -> Self {
29        let mut this = Self::new(s.len());
30
31        for (i, x) in s.iter().enumerate() {
32            this.data[i + this.size / 2] = x.clone();
33        }
34
35        for i in (1..this.size / 2).rev() {
36            this.data[i] = this.data[i << 1]
37                .clone()
38                .op(this.data[(i << 1) | 1].clone());
39        }
40
41        this
42    }
43
44    /// モノイド列をスライスで返す。
45    pub fn to_slice(&self) -> &[M] {
46        &self.data[self.size / 2..self.size / 2 + self.original_size]
47    }
48
49    /// **Time complexity** $O(\log n)$
50    pub fn fold<R: RangeBounds<usize>>(&self, range: R) -> M {
51        let (l, r) = range_bounds_to_range(range, 0, self.size / 2);
52
53        let mut ret_l = M::id();
54        let mut ret_r = M::id();
55
56        let mut l = l + self.size / 2;
57        let mut r = r + self.size / 2;
58
59        while l < r {
60            if r & 1 == 1 {
61                r -= 1;
62                ret_r = M::op(self.data[r].clone(), ret_r);
63            }
64            if l & 1 == 1 {
65                ret_l = M::op(ret_l, self.data[l].clone());
66                l += 1;
67            }
68            r >>= 1;
69            l >>= 1;
70        }
71
72        M::op(ret_l, ret_r)
73    }
74
75    /// **Time complexity** $O(\log n)$
76    pub fn assign(&mut self, i: usize, value: M) {
77        let mut i = i + self.size / 2;
78        self.data[i] = value;
79
80        while i > 1 {
81            i >>= 1;
82            self.data[i] = M::op(self.data[i << 1].clone(), self.data[(i << 1) | 1].clone());
83        }
84    }
85
86    /// **Time complexity** $O(\log n)$
87    pub fn update(&mut self, i: usize, value: M) {
88        self.assign(i, M::op(self.data[i + self.size / 2].clone(), value));
89    }
90}
91
92impl<M: Monoid + Clone> From<&Segtree<M>> for Vec<M> {
93    fn from(from: &Segtree<M>) -> Self {
94        from.data[from.size / 2..from.size / 2 + from.original_size].to_vec()
95    }
96}
97
98impl<M: Monoid> Index<usize> for Segtree<M> {
99    type Output = M;
100
101    fn index(&self, i: usize) -> &Self::Output {
102        &self.data[self.size / 2 + i]
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109    use crate::algebra::traits::*;
110    use my_testtools::*;
111    use rand::Rng;
112
113    fn random_test_helper<M, F>(size: usize, mut gen_value: F)
114    where
115        M: Monoid + Clone + Eq + std::fmt::Debug,
116        F: FnMut() -> M,
117    {
118        let mut rng = rand::thread_rng();
119
120        let mut other = vec![M::id(); size];
121        let mut s = Segtree::new(size);
122
123        for _ in 0..1000 {
124            let ty = rng.gen_range(0..2);
125
126            if ty == 0 {
127                let i = rng.gen_range(0..size);
128                let x = gen_value();
129
130                other[i] = M::op(other[i].clone(), x.clone());
131                s.update(i, x);
132            } else {
133                let lr = rand_range(&mut rng, 0..size);
134
135                let ans = other[lr.clone()].iter().cloned().fold_m();
136
137                assert_eq!(s.fold(lr), ans);
138            }
139
140            let i = rng.gen_range(0..size);
141            assert_eq!(s[i], other[i]);
142        }
143
144        assert_eq!(Vec::<M>::from(&s), other);
145    }
146
147    use crate::algebra::bit::BitXor;
148    use crate::algebra::min_max::{Max, Min};
149    use crate::algebra::sum::Sum;
150
151    #[test]
152    fn test_sum() {
153        let mut rng = rand::thread_rng();
154        random_test_helper(10, || Sum(rng.gen::<i32>() % 10000));
155    }
156
157    #[test]
158    fn test_xor() {
159        let mut rng = rand::thread_rng();
160        random_test_helper(10, || BitXor(rng.gen::<u32>() % 10000));
161    }
162
163    #[test]
164    fn test_min() {
165        let mut rng = rand::thread_rng();
166        random_test_helper(10, || Min(rng.gen::<i32>() % 10000));
167    }
168
169    #[test]
170    fn test_max() {
171        let mut rng = rand::thread_rng();
172        random_test_helper(10, || Max(rng.gen::<i32>() % 10000));
173    }
174}
haar_lib/ds/segtree.rs

haar_lib/ds/
segtree.rs
