competitive_programming_rs/math/
cumulative_sum.rs

1pub struct CumulativeSum<T>(Vec<Vec<T>>);
2
3impl<T> CumulativeSum<T>
4where
5    T: Copy + std::ops::Add<Output = T> + std::ops::Sub<Output = T>,
6{
7    pub fn new(a: &[Vec<T>], init: T) -> CumulativeSum<T> {
8        assert!(!a.is_empty());
9        let h = a.len();
10        let w = a[0].len();
11        let mut sum = vec![vec![init; w + 1]; h + 1];
12        for i in 0..h {
13            for j in 0..w {
14                sum[i + 1][j + 1] = a[i][j] + sum[i][j + 1] + sum[i + 1][j] - sum[i][j];
15            }
16        }
17        CumulativeSum(sum)
18    }
19
20    pub fn get_sum(&self, h_range: std::ops::Range<usize>, w_range: std::ops::Range<usize>) -> T {
21        assert!(h_range.end <= self.0.len());
22        assert!(w_range.end <= self.0[0].len());
23        self.0[h_range.end][w_range.end] + self.0[h_range.start][w_range.start]
24            - self.0[h_range.start][w_range.end]
25            - self.0[h_range.end][w_range.start]
26    }
27}
28
29#[cfg(test)]
30mod test {
31    use super::*;
32    use rand::distributions::Uniform;
33    use rand::Rng;
34
35    #[test]
36    fn random_array() {
37        let h = 30;
38        let w = 20;
39
40        let mut rng = rand::thread_rng();
41
42        for _ in 0..10 {
43            let mut array = vec![vec![0; w]; h];
44            for i in 0..h {
45                for j in 0..w {
46                    array[i][j] = rng.sample(Uniform::from(10..10000));
47                }
48            }
49
50            let sum = CumulativeSum::new(&array, 0);
51            for from_i in 0..h {
52                for to_i in (from_i + 1)..=h {
53                    for from_j in 0..w {
54                        for to_j in (from_j + 1)..=w {
55                            let mut check = 0;
56                            for i in from_i..to_i {
57                                for j in from_j..to_j {
58                                    check += array[i][j];
59                                }
60                            }
61
62                            assert_eq!(check, sum.get_sum(from_i..to_i, from_j..to_j));
63                        }
64                    }
65                }
66            }
67        }
68    }
69}
competitive_programming_rs/math/cumulative_sum.rs

competitive_programming_rs/math/
cumulative_sum.rs
