competitive_library/structure/
fenwick_tree.rs

1//! BIT
2
3pub trait Monoid {
4    type T: Clone;
5    fn identity_element() -> Self::T;
6    fn binary_operation(a: &Self::T, b: &Self::T) -> Self::T;
7}
8
9pub struct Add {}
10impl Monoid for Add {
11    type T = i64;
12    #[inline]
13    fn identity_element() -> Self::T {
14        0_i64
15    }
16    #[inline]
17    fn binary_operation(a: &Self::T, b: &Self::T) -> Self::T {
18        *a + *b
19    }
20}
21
22/// Binary Index Tree
23#[derive(Clone, Debug)]
24pub struct FenwickTree<M>
25where
26    M: Monoid,
27{
28    array: Vec<M::T>,
29}
30
31impl<M> FenwickTree<M>
32where
33    M: Monoid,
34{
35    #[inline]
36    pub fn new(size: usize) -> FenwickTree<M> {
37        Self {
38            array: vec![M::identity_element(); size + 1],
39        }
40    }
41
42    #[inline]
43    pub fn operate(&mut self, index: usize, x: M::T) {
44        let mut i = index + 1;
45        while i < self.array.len() {
46            self.array[i] = M::binary_operation(&self.array[i], &x);
47            i += i & i.wrapping_neg();
48        }
49    }
50
51    /// (0..end)
52    #[inline]
53    pub fn fold(&self, end: usize) -> M::T {
54        let mut s = M::identity_element();
55        let mut i = end;
56        while i > 0 {
57            s = M::binary_operation(&s, &self.array[i]);
58            i -= i & i.wrapping_neg();
59        }
60        s
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67    #[test]
68    fn test_sum() {
69        let mut a = FenwickTree::<Add>::new(100);
70
71        (0..100).for_each(|i| a.operate(i, i as i64 + 1));
72
73        (0..100).for_each(|i| assert_eq!((1..=i).sum::<i64>(), a.fold(i as usize)));
74    }
75
76    pub struct Xor {}
77    impl Monoid for Xor {
78        type T = u64;
79        #[inline]
80        fn identity_element() -> Self::T {
81            0_u64
82        }
83        #[inline]
84        fn binary_operation(a: &Self::T, b: &Self::T) -> Self::T {
85            *a ^ *b
86        }
87    }
88    #[test]
89    fn test_xor() {
90        // https://atcoder.jp/contests/abc185/tasks/abc185_f
91        // sample 2
92        let a = [0, 5, 3, 4, 7, 0, 0, 0, 1, 0];
93        let txy_ans = vec![
94            (1, 10, 7, 0),
95            (2, 8, 9, 1),
96            (2, 3, 6, 0),
97            (2, 1, 6, 5),
98            (2, 1, 10, 3),
99            (1, 9, 4, 0),
100            (1, 6, 1, 0),
101            (1, 6, 3, 0),
102            (1, 1, 7, 0),
103            (2, 3, 5, 0),
104        ];
105
106        let mut ft = FenwickTree::<Xor>::new(10);
107
108        for (i, &v) in a.iter().enumerate() {
109            ft.operate(i, v);
110        }
111
112        for (t, x, y, ans) in txy_ans {
113            if t == 1 {
114                ft.operate(x as usize - 1, y);
115            } else {
116                assert_eq!(ft.fold(y as usize) ^ ft.fold(x as usize - 1), ans);
117            }
118        }
119    }
120}
competitive_library/structure/fenwick_tree.rs

competitive_library/structure/
fenwick_tree.rs
