数列 - Binary Indexed Tree (BIT; Fenwick Tree)

参考

• wikipedia/Fenwick tree

概要

加法群 \((X, +, 0)\) 上の数列について,

• 初期状態
  • 長さ \(n\) の数列: \(v = \{ v_0, v_1, \ldots, v_{n-1} \} = \{0,0,\ldots,0\}\)
• 次の2つの操作が出来る
  1. \(v_i \leftarrow v_i + x\)
  2. \(v_0 + v_1 + \cdots + v_{m-1}\)
     • 特に2つ目から、任意区間の総和が算出できる
• 計算量
  • 時間計算量: 共に \(O(log(n))\)
  • 空間計算量: 数列の長さ \(n\) 程度の配列.

use crate::algebra::group_additive::*;

/// Sequence - Binary Indexed Tree (Fenwick Tree) of Additive Group (+, 0)
pub struct BIT<X> {
    size: usize,
    array: Vec<X>,
}
impl<X: Copy + AGroup> BIT<X> {
    pub fn new(size: usize) -> Self {
        BIT {
            size,
            array: vec![X::zero(); size + 1],
        }
    }
    pub fn add(&mut self, idx: usize, w: X) {
        let mut x = idx + 1;
        while x <= self.size {
            self.array[x] = self.array[x] + w;
            let xi = x as i32;
            x += (xi & -xi) as usize;
        }
    }
    /// sum of [0, idx)
    pub fn sum_up(&self, idx: usize) -> X {
        let mut sum = X::zero();
        let mut x = idx;
        while x > 0 {
            sum = sum + self.array[x];
            let xi = x as i32;
            x -= (xi & -xi) as usize;
        }
        sum
    }
    /// sum of [left, right)
    pub fn sum(&self, range: std::ops::Range<usize>) -> X {
        if range.end <= range.start {
            return X::zero();
        }
        self.sum_up(range.end) - self.sum_up(range.start)
    }
}

#[cfg(test)]
mod test_bit {
    use crate::sequence::tree::bit::*;

    #[test]
    fn sum() {
        let mut bit = BIT::new(10);
        for i in 0..10 {
            bit.add(i, i as i64);
        }
        let mut ac = 0;
        for i in 0..11 {
            assert_eq!(bit.sum_up(i), ac);
            ac = ac + i as i64;
        }
    }
}