haar_lib/ds/
fenwick_add.rs

1//! 可換な加減算に特化したFenwickTree
2use crate::num::one_zero::Zero;
3use crate::trait_alias;
4use std::ops::{Add, Range, RangeTo, Sub};
5
6trait_alias!(
7    /// [`FenwickTreeAdd<T>`]が扱える型
8    Elem: Copy + Zero + Add<Output = Self> + Sub<Output = Self>
9);
10
11/// 可換な加減算に特化したFenwickTree
12pub struct FenwickTreeAdd<T> {
13    data: Vec<T>,
14    size: usize,
15}
16
17impl<T: Elem> FenwickTreeAdd<T> {
18    /// 長さ`size`の[`FenwickTreeAdd<T>`]を生成する。
19    pub fn new(size: usize) -> Self {
20        Self {
21            data: vec![T::zero(); size + 1],
22            size,
23        }
24    }
25
26    /// `i`番目の値から`value`を引く。
27    ///
28    /// **Time complexity** $O(\log n)$
29    pub fn sub(&mut self, mut i: usize, value: T) {
30        i += 1;
31        while i <= self.size {
32            self.data[i] = self.data[i] - value;
33            i += i & (!i + 1);
34        }
35    }
36
37    /// `i`番目の値に`value`を足す。
38    ///
39    /// **Time complexity** $O(\log n)$
40    pub fn add(&mut self, mut i: usize, value: T) {
41        i += 1;
42        while i <= self.size {
43            self.data[i] = self.data[i] + value;
44            i += i & (!i + 1);
45        }
46    }
47
48    /// 範囲`0..r`の総和を返す。
49    ///
50    /// **Time complexity** $O(\log n)$
51    pub fn fold_to(&self, RangeTo { end: mut i }: RangeTo<usize>) -> T {
52        let mut ret = T::zero();
53
54        while i > 0 {
55            ret = ret + self.data[i];
56            i -= i & (!i + 1);
57        }
58
59        ret
60    }
61
62    /// 範囲`l..r`の総和を返す。
63    ///
64    /// **Time complexity** $O(\log n)$
65    pub fn fold(&self, Range { start: l, end: r }: Range<usize>) -> T {
66        self.fold_to(..r) - self.fold_to(..l)
67    }
68}
69
70impl<T: Elem + Ord> FenwickTreeAdd<T> {
71    /// 列の接頭辞の総和が単調増加であるとき、接頭辞の総和が`value`以上となる位置を返す。
72    pub fn lower_bound(&self, value: T) -> usize {
73        let n = self.size;
74        let mut b = 0;
75        let mut len = n;
76
77        while len > 0 {
78            let half = len / 2;
79            let mid = b + half;
80
81            if self.fold_to(..mid + 1) < value {
82                len -= half + 1;
83                b = mid + 1;
84            } else {
85                len = half;
86            }
87        }
88
89        b
90    }
91}
haar_lib/ds/fenwick_add.rs

haar_lib/ds/
fenwick_add.rs
