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 288-295)
278pub fn floor_sum_polynomial<T, const X: usize, const Y: usize>(
279 n: u64,
280 a: u64,
281 b: u64,
282 m: u64,
283) -> [[T; Y]; X]
284where
285 T: Clone + Zero + One + Add<Output = T> + Mul<Output = T>,
286{
287 debug_assert!(a == 0 || n < (u64::MAX - b) / a);
288 floor_monoid_product::<FloorSum<AddMulOperation<T>, X, Y>>(
289 FloorSum::<AddMulOperation<T>, X, Y>::to_x(),
290 FloorSum::<AddMulOperation<T>, X, Y>::to_y(),
291 n,
292 a,
293 b,
294 m,
295 )
296 .dp
297}
298
299/// $$\sum_{i=l}^{r-1}i^X\left\lfloor\frac{a\times i+b}{m}\right\rfloor^Y$$
300pub fn floor_sum_polynomial_i64<T, const X: usize, const Y: usize>(
301 l: i64,
302 r: i64,
303 a: i64,
304 b: i64,
305 m: u64,
306) -> [[T; Y]; X]
307where
308 T: Clone + Zero + One + Add<Output = T> + Mul<Output = T>,
309 AddMulOperation<T>: SemiRing<T = T, Additive: Invertible>,
310{
311 assert!(l <= r);
312 assert!(m > 0);
313
314 if a < 0 {
315 let mut ans = floor_sum_polynomial_i64::<T, X, Y>(-r + 1, -l + 1, -a, b, m);
316 for ans in ans.iter_mut().skip(1).step_by(2) {
317 for ans in ans.iter_mut() {
318 *ans = AddMulOperation::<T>::neg(ans);
319 }
320 }
321 return ans;
322 }
323
324 let add_x = l;
325 let n = (r - l) as u64;
326 let b = a * add_x + b;
327
328 let add_y = b.div_euclid(m as i64);
329 let b = b.rem_euclid(m as i64);
330 assert!(a >= 0);
331 assert!(b >= 0);
332 let data = floor_monoid_product::<FloorSum<AddMulOperation<T>, X, Y>>(
333 FloorSum::<AddMulOperation<T>, X, Y>::to_x(),
334 FloorSum::<AddMulOperation<T>, X, Y>::to_y(),
335 n,
336 a as u64,
337 b as u64,
338 m,
339 );
340
341 let offset = FloorSum::<AddMulOperation<T>, X, Y>::offset(add_x, add_y);
342 FloorSum::<AddMulOperation<T>, X, Y>::operate(&offset, &data).dp
343}
344
345#[derive(Debug)]
346struct FloorPowerSum<R>
347where
348 R: SemiRing,
349{
350 x: R::T,
351 sum: R::T,
352}
353
354impl<R> Clone for FloorPowerSum<R>
355where
356 R: SemiRing,
357{
358 fn clone(&self) -> Self {
359 Self {
360 x: self.x.clone(),
361 sum: self.sum.clone(),
362 }
363 }
364}
365
366impl<R> FloorPowerSum<R>
367where
368 R: SemiRing,
369{
370 fn to_x(x: R::T) -> Self {
371 Self { x, sum: R::one() }
372 }
373 fn to_y(y: R::T) -> Self {
374 Self {
375 x: y,
376 sum: R::zero(),
377 }
378 }
379}
380
381impl<R> Magma for FloorPowerSum<R>
382where
383 R: SemiRing,
384{
385 type T = Self;
386
387 fn operate(a: &Self::T, b: &Self::T) -> Self::T {
388 Self {
389 x: R::mul(&a.x, &b.x),
390 sum: R::add(&a.sum, &R::mul(&a.x, &b.sum)),
391 }
392 }
393}
394
395impl<R> Unital for FloorPowerSum<R>
396where
397 R: SemiRing,
398{
399 fn unit() -> Self::T {
400 Self {
401 x: R::one(),
402 sum: R::zero(),
403 }
404 }
405}
406
407impl<R> Associative for FloorPowerSum<R> where R: SemiRing {}
408
409/// $$\sum_{i=0}^{n-1}x^iy^{\left\lfloor\frac{a\times i+b}{m}\right\rfloor}$$
410pub fn floor_power_sum<R>(x: R::T, y: R::T, n: u64, a: u64, b: u64, m: u64) -> R::T
411where
412 R: SemiRing,
413{
414 floor_monoid_product::<FloorPowerSum<R>>(
415 FloorPowerSum::<R>::to_x(x),
416 FloorPowerSum::<R>::to_y(y),
417 n,
418 a,
419 b,
420 m,
421 )
422 .sum
423}