1use std::ops::{Deref, DerefMut, Index, Range};
2
3use monoid::Monoid;
4use usize_bounds::UsizeBounds;
5
6#[derive(Clone)]
7pub struct VecSegtree<M: Monoid> {
8 tree: Vec<M::Set>,
9 monoid: M,
10}
11
12pub struct PeekMutTmp<'a, M: Monoid> {
13 self_: &'a mut VecSegtree<M>,
14 index: usize,
15}
16
17impl<M: Monoid> VecSegtree<M> {
18 fn init(tree: &mut [M::Set], monoid: &M) {
19 let n = tree.len() / 2;
20 for i in (1..n).rev() {
21 tree[i] = monoid.op(&tree[2 * i], &tree[2 * i + 1]);
22 }
23 }
24 pub fn fold(&self, range: impl UsizeBounds) -> M::Set {
25 let n = self.tree.len() / 2;
26 let monoid = &self.monoid;
27 let Range { start, end } = range.to_range(n);
28 let (mut il, mut ir) = (n + start, n + end);
29 let (mut resl, mut resr) = (monoid.id(), monoid.id());
30 while il < ir {
31 if il & 1 != 0 {
32 resl = monoid.op(&resl, &self.tree[il]);
33 il += 1;
34 }
35 if ir & 1 != 0 {
36 ir -= 1;
37 resr = monoid.op(&self.tree[ir], &resr);
38 }
39 il >>= 1;
40 ir >>= 1;
41 }
42 monoid.op(&resl, &resr)
43 }
44 pub fn peek_mut(&mut self, i: usize) -> PeekMutTmp<'_, M> {
45 let index = self.tree.len() / 2 + i;
46 PeekMutTmp { self_: self, index }
47 }
48
49 fn roots(
50 &self,
51 Range { start, end }: Range<usize>,
52 ) -> impl Iterator<Item = usize> + DoubleEndedIterator {
53 let n = self.tree.len() / 2;
54 let (mut il, mut ir) = (n + start, n + end);
55 let (mut vl, mut vr) = (vec![], vec![]);
56 while il < ir {
57 if il & 1 != 0 {
58 vl.push(il);
59 il += 1;
60 }
61 if ir & 1 != 0 {
62 ir -= 1;
63 vr.push(ir);
64 }
65 il >>= 1;
66 ir >>= 1;
67 }
68 vl.into_iter().chain(vr.into_iter().rev())
69 }
70 pub fn fold_bisect_from<F>(&self, l: usize, pred: F) -> (usize, M::Set)
71 where
72 F: Fn(&M::Set) -> bool,
73 {
74 let n = self.tree.len() / 2;
75 assert!((0..=n).contains(&l));
76
77 let monoid = &self.monoid;
78 let mut x = monoid.id();
79 assert!(pred(&x), "`pred(id)` must hold");
80 match self.fold(l..) {
81 x if pred(&x) => return (n, x),
82 _ => {}
83 }
84
85 for v in self.roots(l..n) {
86 let tmp = monoid.op(&x, &self.tree[v]);
87 if pred(&tmp) {
88 x = tmp;
89 continue;
90 }
91 let mut v = v;
92 while v < n {
93 v *= 2;
94 let tmp = monoid.op(&x, &self.tree[v]);
95 if pred(&tmp) {
96 x = tmp;
97 v += 1;
98 }
99 }
100 return (v - n, x);
101 }
102 unreachable!();
103 }
104 pub fn fold_bisect_to<F>(&self, r: usize, pred: F) -> (usize, M::Set)
105 where
106 F: Fn(&M::Set) -> bool,
107 {
108 let n = self.tree.len() / 2;
109 assert!((0..=n).contains(&r));
110
111 let monoid = &self.monoid;
112 let mut x = monoid.id();
113 assert!(pred(&x), "`pred(id)` must hold");
114 match self.fold(..r) {
115 x if pred(&x) => return (0, x),
116 _ => {}
117 }
118
119 for v in self.roots(0..r).rev() {
120 let tmp = monoid.op(&self.tree[v], &x);
121 if pred(&tmp) {
122 x = tmp;
123 continue;
124 }
125 let mut v = v;
126 while v < n {
127 v = 2 * v + 1;
128 let tmp = monoid.op(&self.tree[v], &x);
129 if pred(&tmp) {
130 x = tmp;
131 v -= 1;
132 }
133 }
134 return (v - n + 1, x);
135 }
136 unreachable!();
137 }
138}
139
140impl<M: Monoid + Default> From<Vec<M::Set>> for VecSegtree<M> {
141 fn from(a: Vec<M::Set>) -> Self {
142 let n = a.len();
143 let monoid = M::default();
144 let mut tree: Vec<_> = (0..n).map(|_| monoid.id()).collect();
145 tree.extend(a);
146 Self::init(&mut tree, &monoid);
147 Self { tree, monoid }
148 }
149}
150
151impl<M: Monoid> From<(Vec<M::Set>, M)> for VecSegtree<M> {
152 fn from((a, monoid): (Vec<M::Set>, M)) -> Self {
153 let n = a.len();
154 let mut tree: Vec<_> = (0..n).map(|_| monoid.id()).collect();
155 tree.extend(a);
156 Self::init(&mut tree, &monoid);
157 Self { tree, monoid }
158 }
159}
160
161impl<M: Monoid> Index<usize> for VecSegtree<M> {
162 type Output = M::Set;
163 fn index(&self, i: usize) -> &Self::Output {
164 let i = self.tree.len() / 2 + i;
165 self.tree.index(i)
166 }
167}
168
169impl<M: Monoid> Deref for PeekMutTmp<'_, M> {
170 type Target = M::Set;
171 fn deref(&self) -> &Self::Target { &self.self_.tree[self.index] }
172}
173
174impl<M: Monoid> DerefMut for PeekMutTmp<'_, M> {
175 fn deref_mut(&mut self) -> &mut Self::Target {
176 &mut self.self_.tree[self.index]
177 }
178}
179
180impl<M: Monoid> Drop for PeekMutTmp<'_, M> {
181 fn drop(&mut self) {
182 let Self {
183 self_: VecSegtree { ref mut tree, ref monoid },
184 index: mut i,
185 } = self;
186 while i > 1 {
187 i >>= 1;
188 tree[i] = monoid.op(&tree[2 * i], &tree[2 * i + 1]);
189 }
190 }
191}
192
193impl<M: Monoid> From<VecSegtree<M>> for Vec<M::Set> {
194 fn from(mut self_: VecSegtree<M>) -> Vec<M::Set> {
195 let n = self_.tree.len() / 2;
196 self_.tree.split_off(n)
197 }
198}
199
200impl<M: Monoid + Default> FromIterator<M::Set> for VecSegtree<M> {
201 fn from_iter<I: IntoIterator<Item = M::Set>>(iter: I) -> Self {
202 let buf: Vec<_> = iter.into_iter().collect();
203 buf.into()
204 }
205}
206
207#[test]
208fn sanity_check() {
209 use op_add::OpAdd;
210
211 let mut tree: VecSegtree<OpAdd<i32>> = vec![1, 2, 3].into();
212 assert_eq!(tree.fold(..), 6);
213 *tree.peek_mut(1) -= 2;
214 assert_eq!(tree.fold(..), 4);
215}
216
217#[test]
218fn fold_bisect() {
219 use op_add::OpAdd;
220
221 let tree: VecSegtree<OpAdd<i32>> = vec![1, 2, 3, 4, 5].into();
222
223 assert_eq!(tree.fold_bisect_from(0, |&x| x <= 0), (0, 0));
224 assert_eq!(tree.fold_bisect_from(0, |&x| x <= 1), (1, 1));
225 assert_eq!(tree.fold_bisect_from(0, |&x| x <= 2), (1, 1));
226 assert_eq!(tree.fold_bisect_from(0, |&x| x <= 3), (2, 3));
227 assert_eq!(tree.fold_bisect_from(0, |&x| x <= 5), (2, 3));
228 assert_eq!(tree.fold_bisect_from(0, |&x| x <= 6), (3, 6));
229 assert_eq!(tree.fold_bisect_from(0, |&x| x <= 9), (3, 6));
230 assert_eq!(tree.fold_bisect_from(0, |&x| x <= 10), (4, 10));
231 assert_eq!(tree.fold_bisect_from(0, |&x| x <= 14), (4, 10));
232 assert_eq!(tree.fold_bisect_from(0, |&x| x <= 15), (5, 15));
233 assert_eq!(tree.fold_bisect_from(0, |_| true), (5, 15));
234
235 assert_eq!(tree.fold_bisect_from(1, |&x| x <= 1), (1, 0));
236 assert_eq!(tree.fold_bisect_from(1, |&x| x <= 2), (2, 2));
237 assert_eq!(tree.fold_bisect_from(1, |&x| x <= 4), (2, 2));
238 assert_eq!(tree.fold_bisect_from(1, |&x| x <= 5), (3, 5));
239 assert_eq!(tree.fold_bisect_from(1, |&x| x <= 8), (3, 5));
240 assert_eq!(tree.fold_bisect_from(1, |&x| x <= 9), (4, 9));
241 assert_eq!(tree.fold_bisect_from(1, |&x| x <= 13), (4, 9));
242 assert_eq!(tree.fold_bisect_from(1, |&x| x <= 14), (5, 14));
243 assert_eq!(tree.fold_bisect_from(1, |_| true), (5, 14));
244
245 assert_eq!(tree.fold_bisect_from(2, |&x| x <= 2), (2, 0));
246 assert_eq!(tree.fold_bisect_from(2, |&x| x <= 3), (3, 3));
247 assert_eq!(tree.fold_bisect_from(2, |&x| x <= 6), (3, 3));
248 assert_eq!(tree.fold_bisect_from(2, |&x| x <= 7), (4, 7));
249 assert_eq!(tree.fold_bisect_from(2, |&x| x <= 11), (4, 7));
250 assert_eq!(tree.fold_bisect_from(2, |&x| x <= 12), (5, 12));
251 assert_eq!(tree.fold_bisect_from(2, |_| true), (5, 12));
252
253 assert_eq!(tree.fold_bisect_from(3, |&x| x <= 3), (3, 0));
254 assert_eq!(tree.fold_bisect_from(3, |&x| x <= 4), (4, 4));
255 assert_eq!(tree.fold_bisect_from(3, |&x| x <= 8), (4, 4));
256 assert_eq!(tree.fold_bisect_from(3, |&x| x <= 9), (5, 9));
257 assert_eq!(tree.fold_bisect_from(3, |_| true), (5, 9));
258
259 assert_eq!(tree.fold_bisect_from(4, |&x| x <= 4), (4, 0));
260 assert_eq!(tree.fold_bisect_from(4, |&x| x <= 5), (5, 5));
261 assert_eq!(tree.fold_bisect_from(4, |_| true), (5, 5));
262
263 assert_eq!(tree.fold_bisect_from(5, |_| true), (5, 0));
264
265 assert_eq!(tree.fold_bisect_to(0, |_| true), (0, 0));
266
267 assert_eq!(tree.fold_bisect_to(1, |&x| x <= 0), (1, 0));
268 assert_eq!(tree.fold_bisect_to(1, |&x| x <= 1), (0, 1));
269 assert_eq!(tree.fold_bisect_to(1, |_| true), (0, 1));
270
271 assert_eq!(tree.fold_bisect_to(2, |&x| x <= 0), (2, 0));
272 assert_eq!(tree.fold_bisect_to(2, |&x| x <= 1), (2, 0));
273 assert_eq!(tree.fold_bisect_to(2, |&x| x <= 2), (1, 2));
274 assert_eq!(tree.fold_bisect_to(2, |&x| x <= 3), (0, 3));
275 assert_eq!(tree.fold_bisect_to(2, |_| true), (0, 3));
276
277 assert_eq!(tree.fold_bisect_to(3, |&x| x <= 2), (3, 0));
278 assert_eq!(tree.fold_bisect_to(3, |&x| x <= 3), (2, 3));
279 assert_eq!(tree.fold_bisect_to(3, |&x| x <= 4), (2, 3));
280 assert_eq!(tree.fold_bisect_to(3, |&x| x <= 5), (1, 5));
281 assert_eq!(tree.fold_bisect_to(3, |&x| x <= 6), (0, 6));
282 assert_eq!(tree.fold_bisect_to(3, |_| true), (0, 6));
283
284 assert_eq!(tree.fold_bisect_to(4, |&x| x <= 0), (4, 0));
285 assert_eq!(tree.fold_bisect_to(4, |&x| x <= 3), (4, 0));
286 assert_eq!(tree.fold_bisect_to(4, |&x| x <= 4), (3, 4));
287 assert_eq!(tree.fold_bisect_to(4, |&x| x <= 6), (3, 4));
288 assert_eq!(tree.fold_bisect_to(4, |&x| x <= 7), (2, 7));
289 assert_eq!(tree.fold_bisect_to(4, |&x| x <= 8), (2, 7));
290 assert_eq!(tree.fold_bisect_to(4, |&x| x <= 9), (1, 9));
291 assert_eq!(tree.fold_bisect_to(4, |&x| x <= 10), (0, 10));
292 assert_eq!(tree.fold_bisect_to(4, |_| true), (0, 10));
293
294 assert_eq!(tree.fold_bisect_to(5, |&x| x <= 0), (5, 0));
295 assert_eq!(tree.fold_bisect_to(5, |&x| x <= 4), (5, 0));
296 assert_eq!(tree.fold_bisect_to(5, |&x| x <= 5), (4, 5));
297 assert_eq!(tree.fold_bisect_to(5, |&x| x <= 8), (4, 5));
298 assert_eq!(tree.fold_bisect_to(5, |&x| x <= 9), (3, 9));
299 assert_eq!(tree.fold_bisect_to(5, |&x| x <= 10), (3, 9));
300 assert_eq!(tree.fold_bisect_to(5, |&x| x <= 12), (2, 12));
301 assert_eq!(tree.fold_bisect_to(5, |&x| x <= 13), (2, 12));
302 assert_eq!(tree.fold_bisect_to(5, |&x| x <= 14), (1, 14));
303 assert_eq!(tree.fold_bisect_to(5, |_| true), (0, 15));
304}