haar_lib/ds/
cumulative_sum_2d.rs

1//! 2次元累積和
2
3pub use crate::algebra::traits::Group;
4use std::ops::{Index, Range};
5
6/// 2次元の累積和を扱う
7#[derive(Debug, Clone)]
8pub struct CumulativeSum2D<G: Group> {
9    data: Vec<Vec<G>>,
10}
11
12/// [`CumulativeSum2D`]を構築する
13pub struct CumulativeSum2DBuilder<G: Group> {
14    data: Vec<Vec<G>>,
15    n: usize,
16    m: usize,
17}
18
19impl<G: Group + Copy> CumulativeSum2D<G> {
20    /// **Time Complexity** $O(1)$
21    pub fn fold_2d(
22        &self,
23        Range { start: l, end: r }: Range<usize>,
24        Range { start: d, end: u }: Range<usize>,
25    ) -> G {
26        let a = self.data[l][u].inv();
27        let b = self.data[r][d].inv();
28        let c = self.data[l][d];
29        let d = self.data[r][u];
30
31        a.op(b.op(c.op(d)))
32    }
33}
34
35impl<G: Group> Index<(usize, usize)> for CumulativeSum2D<G> {
36    type Output = G;
37
38    fn index(&self, (i, j): (usize, usize)) -> &Self::Output {
39        &self.data[i][j]
40    }
41}
42
43impl<G: Group + Copy> CumulativeSum2DBuilder<G> {
44    /// `CumulativeSum2DBuilder`を生成する
45    pub fn new(n: usize, m: usize) -> Self {
46        CumulativeSum2DBuilder {
47            data: vec![vec![G::id(); m + 1]; n + 1],
48            n,
49            m,
50        }
51    }
52
53    /// `[i][j]`番目に`value`を代入する
54    pub fn assign(&mut self, i: usize, j: usize, value: G) {
55        self.data[i + 1][j + 1] = value;
56    }
57
58    /// 群`G`の演算に`[i][j]`番目の値と`value`を適用して`[i][j]`番目の値を更新する。
59    pub fn update(&mut self, i: usize, j: usize, value: G) {
60        self.data[i + 1][j + 1] = self.data[i + 1][j + 1].op(value);
61    }
62
63    /// [`CumulativeSum2D`]を構築する
64    pub fn build(self) -> CumulativeSum2D<G> {
65        let mut data = self.data;
66
67        for i in 1..=self.n {
68            for j in 0..=self.m {
69                data[i][j] = data[i][j].op(data[i - 1][j]);
70            }
71        }
72        for data in data.iter_mut().take(self.n + 1) {
73            for j in 1..=self.m {
74                data[j] = data[j].op(data[j - 1]);
75            }
76        }
77
78        CumulativeSum2D { data }
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85    use crate::algebra::sum::*;
86    use my_testtools::*;
87    use rand::Rng;
88
89    #[test]
90    fn test() {
91        let mut rng = rand::thread_rng();
92
93        let n = 20;
94        let m = 30;
95        let mut csb = CumulativeSum2DBuilder::<Sum<i32>>::new(n, m);
96
97        let mut other = vec![vec![Sum::id(); m]; n];
98
99        for _ in 0..1000 {
100            let i = rng.gen_range(0..n);
101            let j = rng.gen_range(0..m);
102            let x = rng.gen_range(-1000..=1000);
103
104            csb.update(i, j, Sum(x));
105            other[i][j].op_assign_r(Sum(x));
106        }
107
108        let cs = csb.build();
109
110        for _ in 0..100 {
111            let lr = rand_range(&mut rng, 0..n);
112            let du = rand_range(&mut rng, 0..m);
113
114            let ans = other[lr.clone()]
115                .iter()
116                .map(|v| v[du.clone()].iter().cloned().fold_m())
117                .fold_m();
118
119            assert_eq!(cs.fold_2d(lr, du), ans);
120        }
121    }
122}
haar_lib/ds/cumulative_sum_2d.rs

haar_lib/ds/
cumulative_sum_2d.rs
