haar_lib/ds/
fenwick_add.rs

1//! 可換な加減算に特化したFenwickTree
2use crate::num::one_zero::Zero;
3use std::ops::{Add, Range, RangeTo, Sub};
4
5/// 可換な加減算に特化したFenwickTree
6pub struct FenwickTreeAdd<T> {
7    data: Vec<T>,
8    size: usize,
9}
10
11impl<T> FenwickTreeAdd<T>
12where
13    T: Copy + Zero + Add<Output = T> + Sub<Output = T>,
14{
15    /// 長さ`size`の[`FenwickTreeAdd<T>`]を生成する。
16    pub fn new(size: usize) -> Self {
17        Self {
18            data: vec![T::zero(); size + 1],
19            size,
20        }
21    }
22
23    /// `i`番目の値から`value`を引く。
24    ///
25    /// **Time complexity** $O(\log n)$
26    pub fn sub(&mut self, mut i: usize, value: T) {
27        i += 1;
28        while i <= self.size {
29            self.data[i] = self.data[i] - value;
30            i += i & (!i + 1);
31        }
32    }
33
34    /// `i`番目の値に`value`を足す。
35    ///
36    /// **Time complexity** $O(\log n)$
37    pub fn add(&mut self, mut i: usize, value: T) {
38        i += 1;
39        while i <= self.size {
40            self.data[i] = self.data[i] + value;
41            i += i & (!i + 1);
42        }
43    }
44
45    /// 範囲`0..r`の総和を返す。
46    ///
47    /// **Time complexity** $O(\log n)$
48    pub fn fold_to(&self, RangeTo { end: mut i }: RangeTo<usize>) -> T {
49        let mut ret = T::zero();
50
51        while i > 0 {
52            ret = ret + self.data[i];
53            i -= i & (!i + 1);
54        }
55
56        ret
57    }
58
59    /// 範囲`l..r`の総和を返す。
60    ///
61    /// **Time complexity** $O(\log n)$
62    pub fn fold(&self, Range { start: l, end: r }: Range<usize>) -> T {
63        self.fold_to(..r) - self.fold_to(..l)
64    }
65}
66
67impl<T> FenwickTreeAdd<T>
68where
69    T: Copy + Zero + Add<Output = T> + Sub<Output = T> + Ord,
70{
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}