fn floor_monoid_product<M>(
x: M::T,
y: M::T,
n: u64,
a: u64,
b: u64,
m: u64,
) -> M::Twhere
M: Monoid,Examples found in repository?
crates/competitive/src/math/floor_sum.rs (lines 337-344)
327pub fn floor_sum_polynomial<T, const X: usize, const Y: usize>(
328 n: u64,
329 a: u64,
330 b: u64,
331 m: u64,
332) -> [[T; Y]; X]
333where
334 T: DotProduct + One,
335{
336 debug_assert!(a == 0 || n < (u64::MAX - b) / a);
337 floor_monoid_product::<FloorSum<AddMulOperation<T>, X, Y>>(
338 FloorSum::<AddMulOperation<T>, X, Y>::to_x(),
339 FloorSum::<AddMulOperation<T>, X, Y>::to_y(),
340 n,
341 a,
342 b,
343 m,
344 )
345 .dp
346}
347
348/// $$\sum_{i=l}^{r-1}i^X\left\lfloor\frac{a\times i+b}{m}\right\rfloor^Y$$
349pub fn floor_sum_polynomial_i64<T, const X: usize, const Y: usize>(
350 l: i64,
351 r: i64,
352 a: i64,
353 b: i64,
354 m: u64,
355) -> [[T; Y]; X]
356where
357 T: DotProduct + One,
358 AddMulOperation<T>: SemiRing<T = T, Additive: Invertible>,
359{
360 assert!(l <= r);
361 assert!(m > 0);
362
363 if a < 0 {
364 let mut ans = floor_sum_polynomial_i64::<T, X, Y>(-r + 1, -l + 1, -a, b, m);
365 for ans in ans.iter_mut().skip(1).step_by(2) {
366 for ans in ans.iter_mut() {
367 *ans = AddMulOperation::<T>::neg(ans);
368 }
369 }
370 return ans;
371 }
372
373 let add_x = l;
374 let n = (r - l) as u64;
375 let b = a * add_x + b;
376
377 let add_y = b.div_euclid(m as i64);
378 let b = b.rem_euclid(m as i64);
379 assert!(a >= 0);
380 assert!(b >= 0);
381 let data = floor_monoid_product::<FloorSum<AddMulOperation<T>, X, Y>>(
382 FloorSum::<AddMulOperation<T>, X, Y>::to_x(),
383 FloorSum::<AddMulOperation<T>, X, Y>::to_y(),
384 n,
385 a as u64,
386 b as u64,
387 m,
388 );
389
390 let offset = FloorSum::<AddMulOperation<T>, X, Y>::offset(add_x, add_y);
391 FloorSum::<AddMulOperation<T>, X, Y>::operate(&offset, &data).dp
392}
393
394#[derive(Debug)]
395struct FloorPowerSum<R>
396where
397 R: SemiRing,
398{
399 x: R::T,
400 sum: R::T,
401}
402
403impl<R> Clone for FloorPowerSum<R>
404where
405 R: SemiRing,
406{
407 fn clone(&self) -> Self {
408 Self {
409 x: self.x.clone(),
410 sum: self.sum.clone(),
411 }
412 }
413}
414
415impl<R> FloorPowerSum<R>
416where
417 R: SemiRing,
418{
419 fn to_x(x: R::T) -> Self {
420 Self { x, sum: R::one() }
421 }
422 fn to_y(y: R::T) -> Self {
423 Self {
424 x: y,
425 sum: R::zero(),
426 }
427 }
428}
429
430impl<R> Magma for FloorPowerSum<R>
431where
432 R: SemiRing,
433{
434 type T = Self;
435
436 fn operate(a: &Self::T, b: &Self::T) -> Self::T {
437 Self {
438 x: R::mul(&a.x, &b.x),
439 sum: R::add(&a.sum, &R::mul(&a.x, &b.sum)),
440 }
441 }
442}
443
444impl<R> Unital for FloorPowerSum<R>
445where
446 R: SemiRing,
447{
448 fn unit() -> Self::T {
449 Self {
450 x: R::one(),
451 sum: R::zero(),
452 }
453 }
454}
455
456impl<R> Associative for FloorPowerSum<R> where R: SemiRing {}
457
458/// $$\sum_{i=0}^{n-1}x^iy^{\left\lfloor\frac{a\times i+b}{m}\right\rfloor}$$
459pub fn floor_power_sum<R>(x: R::T, y: R::T, n: u64, a: u64, b: u64, m: u64) -> R::T
460where
461 R: SemiRing,
462{
463 floor_monoid_product::<FloorPowerSum<R>>(
464 FloorPowerSum::<R>::to_x(x),
465 FloorPowerSum::<R>::to_y(y),
466 n,
467 a,
468 b,
469 m,
470 )
471 .sum
472}