haar_lib/algo/
imos_2d.rs

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

haar_lib/algo/
imos_2d.rs
