Skip to main content

floor_sum_polynomial_i64

Function floor_sum_polynomial_i64 

Source
pub fn floor_sum_polynomial_i64<T, const X: usize, const Y: usize>(
    l: i64,
    r: i64,
    a: i64,
    b: i64,
    m: u64,
) -> [[T; Y]; X]
where T: DotProduct + One, AddMulOperation<T>: SemiRing<T = T, Additive: Invertible>,
Expand description

$$\sum_{i=l}^{r-1}i^X\left\lfloor\frac{a\times i+b}{m}\right\rfloor^Y$$

Examples found in repository?
crates/competitive/src/math/floor_sum.rs (line 364)
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}