dsalgo/
fenwick_tree_dynamic_cumulative_sum_2_i64.rs

1//! cumulative sum of cumulative sum.
2
3use crate::fenwick_tree_dual_i64_add_1_indexed::Fenwick;
4
5pub struct CumulativeSum2(Fenwick, Fenwick);
6
7impl CumulativeSum2 {
8    pub fn new(size: usize) -> Self {
9        Self(Fenwick::new(size), Fenwick::new(size))
10    }
11
12    pub fn size(&self) -> usize {
13        self.0.size()
14    }
15
16    pub fn add(
17        &mut self,
18        i: usize,
19        x: i64,
20    ) {
21        self.0.add_ge(i, (1 - i as i64) * x);
22
23        self.1.add_ge(i, x);
24    }
25
26    pub fn get(
27        &self,
28        i: usize,
29    ) -> i64 {
30        self.0.get(i) + self.1.get(i) * i as i64
31    }
32}
33
34#[cfg(test)]
35
36mod tests {
37
38    use super::*;
39
40    #[test]
41
42    fn test() {
43        let mut s = CumulativeSum2::new(10);
44
45        s.add(0, 1);
46
47        assert_eq!(s.get(9), 10);
48
49        for i in 1..10 {
50            s.add(i, 1);
51        }
52
53        assert_eq!(s.get(9), 55);
54    }
55}