• + 0 comments

    O((log n)^3) per query/update

    struct FenwickTree {
        tree: Vec<Vec<Vec<i64>>>,
        matr: Vec<Vec<Vec<i32>>>,
        n   : i8
    }
    
    impl FenwickTree {
    
        fn new(n: i8) -> Self {
            Self {
                tree: vec![vec![vec![0; n as usize]; n as usize]; n as usize],
                matr: vec![vec![vec![0; n as usize]; n as usize]; n as usize],
                n
            }
        }
    
        fn up_iter(mut idx: i8, n: i8) -> std::iter::FromFn<impl FnMut() -> Option<usize>> {
            std::iter::from_fn({move || {
                if idx < n {
                    let res = Some(idx as usize);
                    idx |= idx + 1;
                    res
                } else {
                    None
                }
            }})
        }
    
        fn down_iter(mut idx: i8) -> std::iter::FromFn<impl FnMut() -> Option<usize>> {
            std::iter::from_fn({move || {
                if idx >= 0 {
                    let res = Some(idx as usize);
                    idx &= idx + 1;
                    idx -= 1;
                    res
                } else {
                    None
                }
            }})
        }
    
        fn update(&mut self, c: &[i8], w: i64) {
    
            if let [x, y, z] = *c {
                
            let o = &mut self.matr[x as usize][y as usize][z as usize];
    
            let diff = w - *o as i64;
            *o = w as i32;
    
            for i in Self::up_iter(x, self.n) {
                for j in Self::up_iter(y, self.n) {
                    for k in Self::up_iter(z, self.n) {
                        self.tree[i][j][k] += diff;
                    }
                }
            }
            
            } else {panic!("invalid coord input")};
        }
    
        fn query(&self, coords: &mut [i8]) -> i64 {
    
            if let [mut x0, mut y0, mut z0, x1, y1, z1] = *coords {
            
            x0 -= 1; y0 -= 1; z0 -= 1;
    
            let z_update = |x: usize, y: usize| -> i64 {
                let mut res: i64 = 0;
                for z in Self::down_iter(z1) {
                    res += self.tree[x][y][z];
                }
                for z in Self::down_iter(z0) {
                    res -= self.tree[x][y][z];
                }
                res
            };
    
            let y_update = |x: usize| -> i64 {
                let mut res: i64 = 0;
                for y in Self::down_iter(y1) {
                    res += z_update(x, y);
                }
                for y in Self::down_iter(y0) {
                    res -= z_update(x, y);
                }
                res
            };
    
            let x_update = || -> i64 {
                let mut res: i64 = 0;
                for x in Self::down_iter(x1) {
                    res += y_update(x);
                }
                for x in Self::down_iter(x0) {
                    res -= y_update(x);
                }
                res
            };
    
            return x_update();
            
            } else {panic!("invalid coord input")};
        }
    }