haar_lib/ds/
fenwick.rs

1//! 可換群の点更新・区間取得($O(\log n)$, $O(\log n)$)ができる。
2pub use crate::algebra::traits::AbelianGroup;
3use std::ops::{Range, RangeTo};
4
5/// 可換群の点更新・区間取得($O(\log n)$, $O(\log n)$)ができる。
6#[derive(Clone, Default)]
7pub struct FenwickTree<G: AbelianGroup> {
8    data: Vec<G>,
9    size: usize,
10}
11
12impl<G: AbelianGroup + Clone> FenwickTree<G> {
13    /// 長さ`size`、可換群`group`から[`FenwickTree<G>`]を生成する。
14    pub fn new(size: usize) -> Self {
15        let data = vec![G::id(); size + 1];
16        Self { data, size }
17    }
18
19    /// `i`番目の要素を`value`で更新する。
20    ///
21    /// **Time complexity** $O(\log n)$
22    pub fn update(&mut self, mut i: usize, value: G) {
23        i += 1;
24        while i <= self.size {
25            self.data[i] = G::op(self.data[i].clone(), value.clone());
26            i += i & (!i + 1);
27        }
28    }
29
30    /// 範囲`0..r`で計算を集約した結果を返す。
31    ///
32    /// **Time complexity** $O(\log n)$
33    pub fn fold_to(&self, RangeTo { end: mut i }: RangeTo<usize>) -> G {
34        let mut ret = G::id();
35
36        while i > 0 {
37            ret = G::op(ret.clone(), self.data[i].clone());
38            i -= i & (!i + 1);
39        }
40
41        ret
42    }
43
44    /// 範囲`l..r`で計算を集約した結果を返す。
45    ///
46    /// **Time complexity** $O(\log n)$
47    pub fn fold(&self, Range { start: l, end: r }: Range<usize>) -> G {
48        G::op(self.fold_to(..r), self.fold_to(..l).inv())
49    }
50}