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