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

haar_lib/ds/
cumulative_sum_2d.rs
