dsalgo/
fenwick_tree_with_instance_abelian_group_1_indexed.rs

1use crate::fenwick_tree_with_instance_commutative_monoid_1_indexed::*;
2
3pub trait AbelianGroup: Monoid {
4    fn inv(
5        &self,
6        x: Self::T,
7    ) -> Self::T;
8}
9
10impl<G: AbelianGroup> Fenwick<G>
11where
12    G::T: Clone,
13{
14    pub fn fold(
15        &self,
16        l: usize,
17        r: usize,
18    ) -> G::T {
19        assert!(l <= r);
20
21        self.g.op(self.g.inv(self.fold_lt(l)), self.fold_lt(r))
22    }
23
24    pub fn max_right_from<F>(
25        &self,
26        is_ok: F,
27        l: usize,
28    ) -> usize
29    where
30        F: Fn(&G::T) -> bool,
31    {
32        let n = self.size();
33
34        assert!(l <= n);
35
36        let mut i = 0;
37
38        let mut v = self.g.inv(self.fold_lt(l));
39
40        let mut d = (n + 1).next_power_of_two();
41
42        loop {
43            d >>= 1;
44
45            if d == 0 {
46                debug_assert!(l <= i && i <= n);
47
48                return i;
49            }
50
51            if i + d > n {
52                continue;
53            }
54
55            let nv = self.g.op(v.clone(), self.node[i + d].clone());
56
57            if i + d <= l || is_ok(&nv) {
58                i += d;
59
60                v = nv;
61            }
62        }
63    }
64
65    pub fn min_left_from<F>(
66        &self,
67        is_ok: F,
68        r: usize,
69    ) -> usize
70    where
71        F: Fn(&G::T) -> bool,
72    {
73        let n = self.size();
74
75        assert!(r <= n);
76
77        if r == 0 {
78            return 0;
79        }
80
81        let mut v = self.fold_lt(r);
82
83        let mut d = (n + 1).next_power_of_two();
84
85        if is_ok(&v) {
86            return 0;
87        }
88
89        let mut i = 1;
90
91        loop {
92            d >>= 1;
93
94            if d == 0 {
95                debug_assert!(i <= r);
96
97                return i;
98            }
99
100            if i + d > r {
101                continue;
102            }
103
104            let nv =
105                self.g.op(self.g.inv(self.node[i + d - 1].clone()), v.clone());
106
107            if !is_ok(&nv) {
108                i += d;
109
110                v = nv;
111            }
112        }
113    }
114}
115
116#[cfg(test)]
117
118mod tests {
119
120    use super::*;
121
122    #[test]
123
124    fn test() {
125        #[derive(Debug, Clone)]
126
127        struct G;
128
129        impl Monoid for G {
130            type T = i32;
131
132            fn e(&self) -> Self::T {
133                0
134            }
135
136            fn op(
137                &self,
138                l: Self::T,
139                r: Self::T,
140            ) -> Self::T {
141                l + r
142            }
143        }
144
145        impl AbelianGroup for G {
146            fn inv(
147                &self,
148                x: Self::T,
149            ) -> Self::T {
150                -x
151            }
152        }
153
154        let a = vec![0, 1, 0, 0];
155
156        let n = a.len();
157
158        let mut fw = Fenwick::from((G {}, a.as_slice()));
159
160        assert_eq!(fw.fold(0, n), 1);
161
162        fw.operate(2, 1);
163
164        assert_eq!(fw.max_right_from(|v| v <= &1, 1), 2);
165
166        assert_eq!(fw.max_right_from(|v| v <= &1, 2), n);
167
168        assert_eq!(fw.min_left_from(|v| v <= &1, n), 2);
169
170        assert_eq!(fw.min_left_from(|v| v <= &2, n), 0);
171
172        assert_eq!(fw.min_left_from(|v| v <= &1, 2), 0);
173
174        assert_eq!(fw.min_left_from(|v| v <= &0, n), 3);
175    }
176}