区間加算を処理する BIT (Fenwick Tree)

概要

初めゼロからなる数列 \(v = (v_0, v_1, \ldots, v_{n-1})\) について, 次の2つの操作を高速（対数時間）に行いたい.

1. 区間への加算 \(([\ell, r), x)\)
   • \(v_i \leftarrow v_i + x\) for each \(\ell \leq i \lt r\)
2. 要素の取得
   • \(v_i\)

差分列 \(u\) を BIT で管理することにする. つまりある数列 \(u\) があって, これの累積和が元の数列 \(v\) になるようにする.

• \(v_0 = u_0\)
• \(v_1 = u_0 + u_1\)
• \(v_k = \sum_{i=0}^{k} u_i\)

\(v\) の代わりに \(u\) について次のよう処理すればよい. 初期値は \(u_i = 0\) でよくて,

1. 区間への加算 \(([\ell, r), x)\)
   • \(u_\ell \leftarrow u_\ell + x\)
   • \(u_r \leftarrow u_r - x\)
2. 要素の取得
   • \(v_i = \sum_{j=0}^{i} u_j\)

これらは BIT で高速に処理できる.

実装

use crate::algebra::group_additive::*;
use crate::sequence::tree::bit::*;

/// Sequence - Cumulative Array by BIT (Fenwick Tree)
pub struct CumBIT<X> {
    data: BIT<X>,
}
impl<X: Copy + AGroup> CumBIT<X> {
    pub fn new(size: usize) -> Self {
        let data = BIT::new(size);
        Self { data }
    }
    pub fn add(&mut self, range: std::ops::Range<usize>, x: X) {
        self.data.add(range.start, x);
        self.data.add(range.end, -x);
    }
    pub fn at(&self, idx: usize) -> X {
        self.data.sum_up(idx + 1)
    }
}

#[cfg(test)]
mod test_cumbit {
    use crate::sequence::tree::bit_cumulative::*;
    macro_rules! assert_veq {
        ($cbit:expr, $arr:expr) => {
            for i in 0..$arr.len() {
                assert_eq!($cbit.at(i), $arr[i]);
            }
        };
    }
    #[test]
    fn it_works() {
        let mut cbit = CumBIT::new(5);
        assert_veq!(cbit, vec![0, 0, 0, 0, 0]);
        cbit.add(1..4, 2_i128);
        assert_veq!(cbit, vec![0, 2, 2, 2, 0]);
        cbit.add(0..5, 1);
        assert_veq!(cbit, vec![1, 3, 3, 3, 1]);
        cbit.add(3..5, -1);
        assert_veq!(cbit, vec![1, 3, 3, 2, 0]);
        cbit.add(1..2, 100);
        assert_veq!(cbit, vec![1, 103, 3, 2, 0]);
    }
}