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