proconlib/
bit.rs

1use std::ops::{AddAssign, SubAssign};
2
3#[derive(Clone)]
4pub struct BIT<T> {
5    n: usize,
6    bits: Vec<T>,
7}
8 
9impl<T: Default + Copy + AddAssign + SubAssign> BIT<T> {
10    pub fn new(n: usize) -> Self {
11        let mut m = 1;
12        while m < n {
13            m *= 2;
14        }
15        BIT {
16            n: m,
17            bits: vec![T::default(); m + 1],
18        }
19    }
20 
21    pub fn add(&mut self, a: usize, w: T) {
22        let mut x = a as i32 + 1;
23        while x <= self.n as i32 {
24            self.bits[(x as usize)] += w;
25            x += x&-x;
26        }
27    }
28    
29    pub fn sub(&mut self, a: usize, w: T) {
30        let mut x = a as i32 + 1;
31        while x <= self.n as i32 {
32            self.bits[(x as usize)] -= w;
33            x += x&-x;
34        }
35    }
36 
37    pub fn sum(&self, a: usize) -> T {
38        let mut ret = T::default();
39        let mut x = a as i32 + 1;
40        while x > 0 {
41            ret += self.bits[x as usize];
42            x -= x&-x;
43        }
44        ret
45    }
46}
47
48#[cfg(test)]
49mod tests {
50    use super::*;
51
52    #[test]
53    fn bit_test() {
54        let mut bit = BIT::new(8);
55
56        bit.add(1, 1);
57        bit.add(3, 5);
58
59        assert_eq!(0, bit.sum(0));
60        assert_eq!(1, bit.sum(1));
61        assert_eq!(1, bit.sum(2));
62        assert_eq!(6, bit.sum(3));
63        assert_eq!(6, bit.sum(4));
64        assert_eq!(6, bit.sum(7));
65
66        bit.sub(4, 2);
67        assert_eq!(0, bit.sum(0));
68        assert_eq!(1, bit.sum(1));
69        assert_eq!(1, bit.sum(2));
70        assert_eq!(6, bit.sum(3));
71        assert_eq!(4, bit.sum(4));
72        assert_eq!(4, bit.sum(7));
73    }
74}
75
proconlib/bit.rs

proconlib/
bit.rs
