competitive_library/structure/
cumsum_2d.rs

1//! 二次元累積和
2#[derive(Clone, Debug)]
3pub struct CumSum2D {
4    v: Vec<Vec<i64>>,
5}
6
7impl CumSum2D {
8    pub fn new(source: &[Vec<i64>]) -> Self {
9        let h = source.len();
10        let w = source[0].len();
11        let mut v = vec![vec![0i64; w + 1]; h + 1];
12
13        for i in 0..h {
14            for j in 0..w {
15                v[i + 1][j + 1] = source[i][j] + v[i][j + 1] + v[i + 1][j] - v[i][j];
16            }
17        }
18        CumSum2D { v }
19    }
20
21    pub fn query(&self, top: usize, bottom: usize, left: usize, right: usize) -> i64 {
22        self.v[bottom + 1][right + 1] - self.v[bottom + 1][left] - self.v[top][right + 1]
23            + self.v[top][left]
24    }
25}
26
27#[cfg(test)]
28mod tests {
29    use super::*;
30
31    #[test]
32    fn test_cumsum_2d() {
33        let a = CumSum2D::new(&[
34            vec![1, 2, 3, 4],
35            vec![1, 2, 3, 4],
36            vec![1, 2, 3, 4],
37            vec![1, 2, 3, 4],
38        ]);
39        assert_eq!(
40            a.v,
41            vec![
42                vec![0, 0, 0, 0, 0],
43                vec![0, 1, 3, 6, 10],
44                vec![0, 2, 6, 12, 20],
45                vec![0, 3, 9, 18, 30],
46                vec![0, 4, 12, 24, 40]
47            ]
48        );
49        assert_eq!(a.query(0, 0, 0, 0), 1);
50        assert_eq!(a.query(0, 1, 0, 1), 6);
51        assert_eq!(a.query(1, 2, 2, 3), 14);
52        assert_eq!(a.query(0, 0, 0, 3), 10);
53        assert_eq!(a.query(0, 3, 0, 0), 4);
54        assert_eq!(a.query(3, 3, 3, 3), 4);
55
56        assert_eq!(a.query(0, 3, 0, 3), 40);
57    }
58}
competitive_library/structure/cumsum_2d.rs

competitive_library/structure/
cumsum_2d.rs
