1use crate::num::{one_zero::Zero, traits::Signed};
3use std::ops::{Add, Range, Sub};
4
5pub struct Imos1D<T> {
7 data: Vec<T>,
8}
9
10impl<T: Copy + Signed + Zero + Add<Output = T> + Sub<Output = T>> Imos1D<T> {
11 pub fn new(n: usize) -> Self {
13 Self {
14 data: vec![T::zero(); n],
15 }
16 }
17
18 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 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}