haar_lib/algo/
imos_1d.rs

1//! 1次元のimos法
2use crate::num::{one_zero::Zero, traits::Signed};
3use std::ops::{Add, Range, Sub};
4
5/// 1次元のimos法
6pub struct Imos1D<T> {
7    data: Vec<T>,
8}
9
10impl<T: Copy + Signed + Zero + Add<Output = T> + Sub<Output = T>> Imos1D<T> {
11    /// **Time complexity** $O(n)$
12    pub fn new(n: usize) -> Self {
13        Self {
14            data: vec![T::zero(); n],
15        }
16    }
17
18    /// **Time complexity** $O(1)$
19    pub fn update(&mut self, Range { start: l, end: r }: Range<usize>, value: T) {
20        self.data[l] = self.data[l] + value;
21        if let Some(x) = self.data.get_mut(r) {
22            *x = *x - value;
23        }
24    }
25
26    /// **Time complexity** $O(n)$
27    pub fn build(mut self) -> Vec<T> {
28        for i in 1..self.data.len() {
29            self.data[i] = self.data[i] + self.data[i - 1];
30        }
31
32        self.data
33    }
34}
35
36#[cfg(test)]
37mod tests {
38    use super::*;
39    use my_testtools::*;
40    use rand::Rng;
41
42    #[test]
43    fn test() {
44        let n = 100;
45        let t = 1000;
46
47        let mut rng = rand::thread_rng();
48
49        let mut a = Imos1D::<i32>::new(n);
50        let mut ans = vec![0; n];
51
52        for _ in 0..t {
53            let lr = rand_range(&mut rng, 0..n);
54            let x = rng.gen_range(-100..=100);
55
56            a.update(lr.clone(), x);
57
58            for i in lr {
59                ans[i] += x;
60            }
61        }
62
63        assert_eq!(a.build(), ans);
64    }
65}