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