haar_lib/algo/
imos_geo.rs

1//! 等比級数のimos法
2use crate::num::one_zero::{One, Zero};
3use std::ops::{Add, Mul, Range, Sub};
4
5/// 等比級数のimos法
6pub struct ImosGeo<T> {
7    data: Vec<T>,
8    r: T,
9    pow: Vec<T>,
10}
11
12impl<T> ImosGeo<T>
13where
14    T: Copy + Zero + One + Add<Output = T> + Sub<Output = T> + Mul<Output = T>,
15{
16    /// **Time complexity** $O(n)$
17    pub fn new(n: usize, r: T) -> Self {
18        let mut pow = vec![T::one(); n];
19        for i in 1..n {
20            pow[i] = pow[i - 1] * r;
21        }
22
23        Self {
24            data: vec![T::zero(); n],
25            r,
26            pow,
27        }
28    }
29
30    /// **Time complexity** $O(1)$
31    pub fn update(&mut self, Range { start: l, end: r }: Range<usize>, value: T) {
32        self.data[l] = self.data[l] + value;
33        if let Some(x) = self.data.get_mut(r) {
34            *x = *x - self.pow[r - l] * value;
35        }
36    }
37
38    /// **Time complexity** $O(n)$
39    pub fn build(mut self) -> Vec<T> {
40        for i in 1..self.data.len() {
41            self.data[i] = self.data[i] + self.data[i - 1] * self.r;
42        }
43
44        self.data
45    }
46}
47
48#[cfg(test)]
49mod tests {
50    use crate::num::const_modint::*;
51    use my_testtools::rand_range;
52
53    use super::*;
54    use rand::Rng;
55
56    #[test]
57    fn test() {
58        let mut rng = rand::thread_rng();
59
60        let ff = ConstModIntBuilder::<1000000007>;
61
62        let n = 100;
63        let t = 100;
64
65        for r in 2..=100 {
66            let r = ff.from_u64(r);
67
68            let mut a = ImosGeo::new(n, r);
69            let mut ans = vec![ff.from_u64(0); n];
70
71            for _ in 0..t {
72                let lr = rand_range(&mut rng, 0..n);
73                let x = ff.from_i64(rng.gen_range(-100..=100));
74
75                a.update(lr.clone(), x);
76
77                let mut x = x;
78                for i in lr {
79                    ans[i] += x;
80                    x *= r;
81                }
82            }
83
84            assert_eq!(a.build(), ans);
85        }
86    }
87}
haar_lib/algo/imos_geo.rs

haar_lib/algo/
imos_geo.rs
