1use super::{AbelianGroup, AbelianMonoid, Group, Monoid, RangeBoundsExt};
2use std::{
3 fmt::{self, Debug, Formatter},
4 iter::FromIterator,
5 ops::RangeBounds,
6};
7
8pub struct Accumulate<M>
10where
11 M: Monoid,
12{
13 data: Vec<M::T>,
14}
15
16impl<M> Debug for Accumulate<M>
17where
18 M: Monoid,
19 M::T: Debug,
20{
21 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
22 f.debug_struct("Accumulate")
23 .field("data", &self.data)
24 .finish()
25 }
26}
27
28impl<M> FromIterator<M::T> for Accumulate<M>
29where
30 M: Monoid,
31{
32 fn from_iter<T>(iter: T) -> Self
33 where
34 T: IntoIterator<Item = M::T>,
35 {
36 let iter = iter.into_iter();
37 let (lower, _) = iter.size_hint();
38 let mut data = Vec::with_capacity(lower.saturating_add(1));
39 let mut acc = M::unit();
40 for x in iter {
41 let y = M::operate(&acc, &x);
42 data.push(acc);
43 acc = y;
44 }
45 data.push(acc);
46 Self { data }
47 }
48}
49
50impl<M> Accumulate<M>
51where
52 M: Monoid,
53{
54 pub fn accumulate(&self, k: usize) -> M::T {
56 assert!(
57 k < self.data.len(),
58 "index out of range: the len is {} but the index is {}",
59 self.data.len(),
60 k
61 );
62 unsafe { self.data.get_unchecked(k) }.clone()
63 }
64}
65
66impl<M> Accumulate<M>
67where
68 M: Group,
69{
70 pub fn fold<R>(&self, range: R) -> M::T
72 where
73 R: RangeBounds<usize>,
74 {
75 let n = self.data.len() - 1;
76 let range = range.to_range_bounded(0, n).expect("invalid range");
77 let (l, r) = (range.start, range.end);
78 assert!(l <= r, "bad range [{}, {})", l, r);
79 M::operate(&M::inverse(unsafe { self.data.get_unchecked(l) }), unsafe {
80 self.data.get_unchecked(r)
81 })
82 }
83}
84
85pub struct Accumulate2d<M>
87where
88 M: AbelianMonoid,
89{
90 h: usize,
91 w: usize,
92 data: Vec<M::T>,
93}
94
95impl<M> Debug for Accumulate2d<M>
96where
97 M: AbelianMonoid,
98 M::T: Debug,
99{
100 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
101 f.debug_struct("Accumulate2d")
102 .field("h", &self.h)
103 .field("w", &self.w)
104 .field("data", &self.data)
105 .finish()
106 }
107}
108
109impl<M> Accumulate2d<M>
110where
111 M: AbelianMonoid,
112{
113 pub fn new(arr2d: &[Vec<M::T>]) -> Self {
114 let h = arr2d.len();
115 assert!(h > 0);
116 let w = arr2d[0].len();
117 assert!(w > 0);
118 let w1 = w + 1;
119 let mut data = Vec::with_capacity((h + 1) * w1);
120 data.resize_with(w1, M::unit);
121 for (i, arr) in arr2d.iter().enumerate() {
122 assert_eq!(w, arr.len(), "expected 2d array");
123 let mut acc = M::unit();
124 for (j, x) in arr.iter().enumerate() {
125 let y = M::operate(&acc, x);
126 data.push(M::operate(&acc, unsafe { data.get_unchecked(w1 * i + j) }));
127 acc = y;
128 }
129 data.push(M::operate(&acc, unsafe { data.get_unchecked(w1 * i + w) }));
130 }
131 Self { h, w, data }
132 }
133 pub fn from_fn<F>(h: usize, w: usize, mut f: F) -> Self
134 where
135 F: FnMut(usize, usize) -> M::T,
136 {
137 let w1 = w + 1;
138 let mut data = Vec::with_capacity((h + 1) * w1);
139 data.resize_with(w1, M::unit);
140 for i in 0..h {
141 let mut acc = M::unit();
142 for j in 0..w {
143 let y = M::operate(&acc, &f(i, j));
144 data.push(M::operate(&acc, unsafe { data.get_unchecked(w1 * i + j) }));
145 acc = y;
146 }
147 data.push(M::operate(&acc, unsafe { data.get_unchecked(w1 * i + w) }));
148 }
149 Self { h, w, data }
150 }
151 pub fn accumulate(&self, x: usize, y: usize) -> M::T {
153 let h1 = self.h + 1;
154 let w1 = self.w + 1;
155 assert!(
156 x < h1,
157 "index out of range: the first len is {} but the index is {}",
158 h1,
159 x
160 );
161 assert!(
162 y < w1,
163 "index out of range: the second len is {} but the index is {}",
164 w1,
165 y
166 );
167 unsafe { self.data.get_unchecked(w1 * x + y) }.clone()
168 }
169}
170
171impl<M> Accumulate2d<M>
172where
173 M: AbelianGroup,
174{
175 pub fn fold<R0, R1>(&self, range0: R0, range1: R1) -> M::T
177 where
178 R0: RangeBounds<usize>,
179 R1: RangeBounds<usize>,
180 {
181 let range0 = range0.to_range_bounded(0, self.h).expect("invalid range");
182 let range1 = range1.to_range_bounded(0, self.w).expect("invalid range");
183 let (xl, xr) = (range0.start, range0.end);
184 let (yl, yr) = (range1.start, range1.end);
185 assert!(xl <= xr, "bad range [{}, {})", xl, xr);
186 assert!(yl <= yr, "bad range [{}, {})", yl, yr);
187 let w1 = self.w + 1;
188 unsafe {
189 M::rinv_operate(
190 &M::operate(
191 self.data.get_unchecked(w1 * xl + yl),
192 self.data.get_unchecked(w1 * xr + yr),
193 ),
194 &M::operate(
195 self.data.get_unchecked(w1 * xl + yr),
196 self.data.get_unchecked(w1 * xr + yl),
197 ),
198 )
199 }
200 }
201}
202
203pub struct AccumulateKd<const K: usize, M>
204where
205 M: AbelianMonoid,
206{
207 dim: [usize; K],
208 offset: [usize; K],
209 data: Vec<M::T>,
210}
211
212impl<const K: usize, M> Debug for AccumulateKd<K, M>
213where
214 M: AbelianMonoid,
215 M::T: Debug,
216{
217 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
218 f.debug_struct("AccumulateKd")
219 .field("dim", &self.dim)
220 .field("offset", &self.offset)
221 .field("data", &self.data)
222 .finish()
223 }
224}
225
226impl<const K: usize, M> AccumulateKd<K, M>
227where
228 M: AbelianMonoid,
229{
230 pub fn from_fn(dim: [usize; K], mut f: impl FnMut([usize; K]) -> M::T) -> Self {
231 fn fill<const K: usize, T>(
232 dim: &[usize; K],
233 offset: &[usize; K],
234 data: &mut [T],
235 f: &mut impl FnMut([usize; K]) -> T,
236 mut index: [usize; K],
237 pos: usize,
238 ) {
239 if pos < K {
240 for i in 0..dim[pos] {
241 index[pos] = i;
242 fill(dim, offset, data, f, index, pos + 1);
243 }
244 } else {
245 let i: usize = index.iter().zip(offset).map(|(x, y)| (x + 1) * y).sum();
246 data[i] = f(index);
247 }
248 }
249
250 let mut offset = [1; K];
251 for d in (1..K).rev() {
252 offset[d - 1] = offset[d] * (dim[d] + 1);
253 }
254 let size = offset[0] * (dim[0] + 1);
255 let mut data = vec![M::unit(); size];
256 fill(&dim, &offset, &mut data, &mut f, [0; K], 0);
257 for d in 0..K {
258 for i in 1..size {
259 if i / offset[d] % (dim[d] + 1) != 0 {
260 data[i] = M::operate(&data[i], &data[i - offset[d]]);
261 }
262 }
263 }
264 Self { dim, offset, data }
265 }
266 pub fn accumulate(&self, x: [usize; K]) -> M::T {
267 for (d, x) in x.into_iter().enumerate() {
268 assert!(
269 x <= self.dim[d],
270 "index out of range: the len is {} but the index is {}",
271 self.dim[d] + 1,
272 x
273 );
274 }
275 let p: usize = x.iter().zip(&self.offset).map(|(x, y)| x * y).sum();
276 unsafe { self.data.get_unchecked(p) }.clone()
277 }
278}
279
280impl<const K: usize, M> AccumulateKd<K, M>
281where
282 M: AbelianGroup,
283{
284 pub fn fold<R>(&self, ranges: [R; K]) -> M::T
285 where
286 R: RangeBounds<usize>,
287 {
288 let ranges: [_; K] = std::array::from_fn(|i| {
289 let range = ranges[i]
290 .to_range_bounded(0, self.dim[i])
291 .expect("invalid range");
292 let (l, r) = (range.start, range.end);
293 assert!(l <= r, "bad range [{}, {})", l, r);
294 [l, r]
295 });
296 let mut acc = M::unit();
297 for bit in 0..1 << K {
298 let p: usize = ranges
299 .iter()
300 .zip(&self.offset)
301 .enumerate()
302 .map(|(d, (range, offset))| range[(bit >> d) & 1 ^ 1] * offset)
303 .sum();
304 if bit.count_ones() & 1 == 0 {
305 acc = M::operate(&acc, unsafe { self.data.get_unchecked(p) });
306 } else {
307 acc = M::rinv_operate(&acc, unsafe { self.data.get_unchecked(p) });
308 }
309 }
310 acc
311 }
312}
313
314#[cfg(test)]
315mod tests {
316 use super::*;
317 use crate::{
318 algebra::{AdditiveOperation, LinearOperation, Magma, Unital},
319 num::mint_basic::MInt1000000007,
320 rand,
321 tools::{RandomSpec, Xorshift},
322 };
323 type M = LinearOperation<MInt1000000007>;
324 type A = AdditiveOperation<MInt1000000007>;
325 struct D;
326 impl RandomSpec<MInt1000000007> for D {
327 fn rand(&self, rng: &mut Xorshift) -> MInt1000000007 {
328 MInt1000000007::new_unchecked(rng.random(..MInt1000000007::get_mod()))
329 }
330 }
331
332 #[test]
333 fn test_accumlate() {
334 let mut rng = Xorshift::default();
335 const Q: usize = 1_000;
336 const N: usize = 50;
337 for n in 0..Q {
338 let n = n % N;
339 rand!(rng, v: [(D, D); n]);
340 let acc: Accumulate<M> = v.iter().cloned().collect();
341 for r in 0..=n {
342 assert_eq!(
343 v[..r].iter().fold(M::unit(), |x, y| M::operate(&x, y)),
344 acc.accumulate(r)
345 );
346 for l in 0..=r {
347 assert_eq!(
348 v[l..r].iter().fold(M::unit(), |x, y| M::operate(&x, y)),
349 acc.fold(l..r)
350 );
351 }
352 }
353 }
354 }
355
356 #[test]
357 fn test_accumlate2d() {
358 let mut rng = Xorshift::default();
359 const Q: usize = 1_000;
360 const N: usize = 10;
361 for i in 0..Q {
362 let h = i % N + 1;
363 let w = i / N % N + 1;
364 rand!(rng, v: [[D; w]; h]);
365 let acc2d = Accumulate2d::<A>::new(&v);
366 for xr in 0..=h {
367 for yr in 0..=w {
368 assert_eq!(
369 v[..xr]
370 .iter()
371 .flat_map(|v| v[..yr].iter())
372 .fold(A::unit(), |x, y| A::operate(&x, y)),
373 acc2d.accumulate(xr, yr)
374 );
375 for xl in 0..=xr {
376 for yl in 0..=yr {
377 assert_eq!(
378 v[xl..xr]
379 .iter()
380 .flat_map(|v| v[yl..yr].iter())
381 .fold(A::unit(), |x, y| A::operate(&x, y)),
382 acc2d.fold(xl..xr, yl..yr)
383 );
384 }
385 }
386 }
387 }
388 }
389 }
390
391 #[test]
392 fn test_accumlate2d_from_fn() {
393 let mut rng = Xorshift::default();
394 const Q: usize = 1_000;
395 const N: usize = 10;
396 for i in 0..Q {
397 let h = i % N;
398 let w = i / N % N;
399 rand!(rng, v: [[D; w]; h]);
400 let acc2d = Accumulate2d::<A>::from_fn(h, w, |i, j| v[i][j]);
401 for xr in 0..=h {
402 for yr in 0..=w {
403 assert_eq!(
404 v[..xr]
405 .iter()
406 .flat_map(|v| v[..yr].iter())
407 .fold(A::unit(), |x, y| A::operate(&x, y)),
408 acc2d.accumulate(xr, yr)
409 );
410 for xl in 0..=xr {
411 for yl in 0..=yr {
412 assert_eq!(
413 v[xl..xr]
414 .iter()
415 .flat_map(|v| v[yl..yr].iter())
416 .fold(A::unit(), |x, y| A::operate(&x, y)),
417 acc2d.fold(xl..xr, yl..yr)
418 );
419 }
420 }
421 }
422 }
423 }
424 }
425
426 #[test]
427 fn test_accumlatekd_from_fn_3d() {
428 let mut rng = Xorshift::default();
429 const N: usize = 5;
430 for i in 0..N * N * N {
431 let dim = [i % N, i / N % N, i / N / N % N];
432 rand!(rng, v: [[[D; dim[2]]; dim[1]]; dim[0]]);
433 let acc = AccumulateKd::<3, A>::from_fn(dim, |[i, j, k]| v[i][j][k]);
434 for xr in 0..=dim[0] {
435 for yr in 0..=dim[1] {
436 for zr in 0..=dim[2] {
437 assert_eq!(
438 v[..xr]
439 .iter()
440 .flat_map(|v| v[..yr].iter().flat_map(|v| v[..zr].iter()))
441 .fold(A::unit(), |x, y| A::operate(&x, y)),
442 acc.accumulate([xr, yr, zr])
443 );
444 for xl in 0..=xr {
445 for yl in 0..=yr {
446 for zl in 0..=zr {
447 assert_eq!(
448 v[xl..xr]
449 .iter()
450 .flat_map(|v| v[yl..yr]
451 .iter()
452 .flat_map(|v| v[zl..zr].iter()))
453 .fold(A::unit(), |x, y| A::operate(&x, y)),
454 acc.fold([xl..xr, yl..yr, zl..zr])
455 );
456 }
457 }
458 }
459 }
460 }
461 }
462 }
463 }
464
465 #[test]
466 fn test_accumlatekd_from_fn_4d() {
467 let mut rng = Xorshift::default();
468 const N: usize = 4;
469 for i in 0..N * N * N * N {
470 let dim = [i % N, i / N % N, i / N / N % N, i / N / N / N % N];
471 rand!(rng, v: [[[[D; dim[3]]; dim[2]]; dim[1]]; dim[0]]);
472 let acc = AccumulateKd::<4, A>::from_fn(dim, |[i, j, k, l]| v[i][j][k][l]);
473 for xr in 0..=dim[0] {
474 for yr in 0..=dim[1] {
475 for zr in 0..=dim[2] {
476 for wr in 0..=dim[3] {
477 assert_eq!(
478 v[..xr]
479 .iter()
480 .flat_map(|v| v[..yr]
481 .iter()
482 .flat_map(|v| v[..zr].iter().flat_map(|v| v[..wr].iter())))
483 .fold(A::unit(), |x, y| A::operate(&x, y)),
484 acc.accumulate([xr, yr, zr, wr])
485 );
486 for xl in 0..=xr {
487 for yl in 0..=yr {
488 for zl in 0..=zr {
489 for wl in 0..=wr {
490 assert_eq!(
491 v[xl..xr]
492 .iter()
493 .flat_map(|v| v[yl..yr]
494 .iter()
495 .flat_map(|v| v[zl..zr]
496 .iter()
497 .flat_map(|v| v[wl..wr].iter())))
498 .fold(A::unit(), |x, y| A::operate(&x, y)),
499 acc.fold([xl..xr, yl..yr, zl..zr, wl..wr])
500 );
501 }
502 }
503 }
504 }
505 }
506 }
507 }
508 }
509 }
510 }
511}