competitive/math/
floor_sum.rs

1use super::{
2    AddMulOperation, Associative, BarrettReduction, DotProduct, Group, Invertible, Magma, Monoid,
3    One, Ring, SemiRing, Unital, Wrapping, array,
4};
5use std::{marker::PhantomData, mem::swap, ops::Range};
6
7fn choose2(n: Wrapping<u64>) -> Wrapping<u64> {
8    if n.0.is_multiple_of(2) {
9        n / 2 * (n - 1)
10    } else {
11        (n - 1) / 2 * n
12    }
13}
14
15/// Sum of Floor of Linear mod 2^64
16///
17/// $$\sum_{i=0}^{n-1}\left\lfloor\frac{a\times i+b}{m}\right\rfloor$$
18pub fn floor_sum(n: u64, a: u64, b: u64, m: u64) -> u64 {
19    let mut ans = Wrapping(0u64);
20    let (mut n, mut m, mut a, mut b) = (Wrapping(n), m, a, b);
21    loop {
22        let br = BarrettReduction::<u64>::new(m);
23        if a >= m {
24            let (q, r) = br.div_rem(a);
25            ans += choose2(n) * q;
26            a = r;
27        }
28        if b >= m {
29            let (q, r) = br.div_rem(b);
30            ans += n * q;
31            b = r;
32        }
33        let y_max = (n * a + b).0;
34        if y_max < m {
35            break;
36        }
37        let (q, r) = br.div_rem(y_max);
38        n = Wrapping(q);
39        b = r;
40        swap(&mut m, &mut a);
41    }
42    ans.0
43}
44
45/// $$\min(\{(a\times i+b)\bmod m\mid0\leq i<n\}\cup\{m\})$$
46pub fn min_of_mod_of_linear(mut n: u64, mut a: u64, mut b: u64, mut m: u64) -> u64 {
47    if n == 0 {
48        return m;
49    }
50    if a >= m {
51        a %= m;
52    }
53    if b >= m {
54        b %= m;
55    }
56    let mut ans = Wrapping(0u64);
57    let mut pos = true;
58    let mut p = 1;
59    let mut q = 1;
60    while a != 0 {
61        let e = if pos { a } else { m } - 1;
62        let r = m % a;
63        let d = m - b;
64        if if pos { b + 1 } else { d } > a {
65            let t = (d - 1) / a + u64::from(pos);
66            let c = (t - u64::from(pos)) * p + if pos { q } else { 0 };
67            if n <= c {
68                if !pos {
69                    ans -= Wrapping(a) * Wrapping((n - 1) / p);
70                }
71                break;
72            }
73            n -= c;
74            if pos {
75                b = a * t - d;
76            } else {
77                b += a * t;
78            }
79        }
80        if r != 0 {
81            let x = m / a * p + q;
82            q = x;
83            p = x - p;
84        }
85        if pos {
86            ans += e;
87        } else {
88            ans -= e;
89        }
90        m = a;
91        a = r;
92        b = e - b;
93        pos = !pos;
94    }
95    if pos { (ans + b).0 } else { (ans - b).0 }
96}
97
98/// Sum of Floor of Linear mod 2^64
99///
100/// $$\sum_{i=l}^{r-1}\left\lfloor\frac{a\times i+b}{m}\right\rfloor$$
101pub fn floor_sum_i64(l: i64, r: i64, a: i64, b: i64, m: u64) -> i64 {
102    let mut ans = Wrapping(0i64);
103    let (n, m, a, b) = (
104        Wrapping((r - l) as u64),
105        m as i64,
106        a,
107        (Wrapping(l) * a + b).0,
108    );
109    let a = if a < 0 {
110        let r = a.rem_euclid(m);
111        let nc2 = choose2(n);
112        ans -= Wrapping(nc2.0 as _) * ((Wrapping(r) - a) / m);
113        r
114    } else {
115        a
116    };
117    let b = if b < 0 {
118        let r = b.rem_euclid(m);
119        ans -= Wrapping(n.0 as _) * ((Wrapping(r) - b) / m);
120        r
121    } else {
122        b
123    };
124    (ans + floor_sum(n.0, a as u64, b as u64, m as u64) as i64).0
125}
126
127pub fn floor_sum_range_freq(l: i64, r: i64, a: i64, b: i64, m: u64, range: Range<i64>) -> i64 {
128    if range.start >= range.end {
129        return 0;
130    }
131    assert!(0 <= range.start && range.end <= m as i64);
132    let x1 = floor_sum_i64(l, r, a, b - range.start, m);
133    let x2 = floor_sum_i64(l, r, a, b - range.end, m);
134    x1 - x2
135}
136
137struct FloorSum<R, const X: usize, const Y: usize>
138where
139    R: SemiRing,
140{
141    _marker: PhantomData<fn() -> R>,
142}
143
144#[derive(Debug)]
145struct FloorSumData<R, const X: usize, const Y: usize>
146where
147    R: SemiRing,
148{
149    dp: [[R::T; Y]; X],
150    dx: R::T,
151    dy: R::T,
152    _marker: PhantomData<fn() -> R>,
153}
154
155impl<R, const X: usize, const Y: usize> Clone for FloorSumData<R, X, Y>
156where
157    R: SemiRing,
158{
159    fn clone(&self) -> Self {
160        Self {
161            dp: self.dp.clone(),
162            dx: self.dx.clone(),
163            dy: self.dy.clone(),
164            _marker: self._marker,
165        }
166    }
167}
168
169impl<R, const X: usize, const Y: usize> FloorSum<R, X, Y>
170where
171    R: SemiRing,
172{
173    fn to_x() -> FloorSumData<R, X, Y> {
174        let mut dp = array![array![R::zero(); Y]; X];
175        dp[0][0] = R::one();
176        FloorSumData {
177            dp,
178            dx: R::one(),
179            dy: R::zero(),
180            _marker: PhantomData,
181        }
182    }
183    fn to_y() -> FloorSumData<R, X, Y> {
184        FloorSumData {
185            dp: array![array![R::zero(); Y]; X],
186            dx: R::zero(),
187            dy: R::one(),
188            _marker: PhantomData,
189        }
190    }
191}
192
193impl<R, const X: usize, const Y: usize> FloorSum<R, X, Y>
194where
195    R: Ring<Additive: Invertible>,
196{
197    fn offset(x: i64, y: i64) -> FloorSumData<R, X, Y> {
198        FloorSumData {
199            dp: array![array![R::zero(); Y]; X],
200            dx: R::Additive::signed_pow(R::one(), x),
201            dy: R::Additive::signed_pow(R::one(), y),
202            _marker: PhantomData,
203        }
204    }
205}
206
207impl<R, const X: usize, const Y: usize> Magma for FloorSum<R, X, Y>
208where
209    R: SemiRing,
210{
211    type T = FloorSumData<R, X, Y>;
212
213    fn operate(a: &Self::T, b: &Self::T) -> Self::T {
214        let mut a = a.clone();
215        let mut b = b.clone();
216        let mut pow_x = array![R::zero(); X];
217        let mut pow_y = array![R::zero(); Y];
218        pow_x[0] = R::one();
219        pow_y[0] = R::one();
220        for i in 1..X {
221            pow_x[i] = R::mul(&pow_x[i - 1], &a.dx);
222        }
223        for j in 1..Y {
224            pow_y[j] = R::mul(&pow_y[j - 1], &a.dy);
225        }
226        macro_rules! go {
227            ($N:ident) => {
228                let mut comb = array![array![R::zero(); $N]; $N];
229                comb[0][0] = R::one();
230                let mut i = 0;
231                while i + 1 < $N {
232                    let mut j = 0;
233                    while j <= i {
234                        let x = comb[i][j].clone();
235                        R::add_assign(&mut comb[i + 1][j], &x);
236                        R::add_assign(&mut comb[i + 1][j + 1], &x);
237                        j += 1;
238                    }
239                    i += 1;
240                }
241                for i in 0..X {
242                    for j in (0..Y).rev() {
243                        for k in j + 1..Y {
244                            let mut x = b.dp[i][j].clone();
245                            R::mul_assign(&mut x, &comb[k][j]);
246                            R::mul_assign(&mut x, &pow_y[k - j]);
247                            R::add_assign(&mut b.dp[i][k], &x);
248                        }
249                    }
250                }
251                for j in 0..Y {
252                    for i in (0..X).rev() {
253                        for k in i..X {
254                            let mut x = b.dp[i][j].clone();
255                            R::mul_assign(&mut x, &comb[k][i]);
256                            R::mul_assign(&mut x, &pow_x[k - i]);
257                            R::add_assign(&mut a.dp[k][j], &x);
258                        }
259                    }
260                }
261            };
262        }
263        if X <= Y {
264            go!(Y);
265        } else {
266            go!(X);
267        }
268        R::add_assign(&mut a.dx, &b.dx);
269        R::add_assign(&mut a.dy, &b.dy);
270        a
271    }
272}
273
274impl<R, const X: usize, const Y: usize> Unital for FloorSum<R, X, Y>
275where
276    R: SemiRing,
277{
278    fn unit() -> Self::T {
279        FloorSumData {
280            dp: array![array![R::zero(); Y]; X],
281            dx: R::zero(),
282            dy: R::zero(),
283            _marker: PhantomData,
284        }
285    }
286}
287
288impl<R, const X: usize, const Y: usize> Associative for FloorSum<R, X, Y> where R: SemiRing {}
289
290fn floor_monoid_product<M>(
291    mut x: M::T,
292    mut y: M::T,
293    mut n: u64,
294    mut a: u64,
295    mut b: u64,
296    mut m: u64,
297) -> M::T
298where
299    M: Monoid,
300{
301    let mut c = (a * n + b) / m;
302    let mut pre = M::unit();
303    let mut suf = M::unit();
304    loop {
305        let (p, q) = (a / m, b / m);
306        a %= m;
307        b %= m;
308        x = M::operate(&x, &M::pow(y.clone(), p));
309        pre = M::operate(&pre, &M::pow(y.clone(), q));
310        c -= p * n + q;
311        if c == 0 {
312            break;
313        }
314        let d = (m * c - b - 1) / a + 1;
315        suf = M::operate(&y, &M::operate(&M::pow(x.clone(), n - d), &suf));
316        b = m - b - 1 + a;
317        n = c - 1;
318        c = d;
319        swap(&mut m, &mut a);
320        swap(&mut x, &mut y);
321    }
322    x = M::pow(x.clone(), n);
323    M::operate(&M::operate(&pre, &x), &suf)
324}
325
326/// $$\sum_{i=0}^{n-1}i^X\left\lfloor\frac{a\times i+b}{m}\right\rfloor^Y$$
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}
473
474#[cfg(test)]
475mod tests {
476    use super::*;
477    use crate::{num::mint_basic::MInt998244353, tools::Xorshift};
478
479    #[test]
480    fn test_floor_sum() {
481        const A: u64 = 1_000;
482        const B: i64 = 1_000;
483        const Q: usize = 1_000;
484        let mut rng = Xorshift::default();
485        for _ in 0..Q {
486            let (n, a, b, m) = rng.random((..A, ..A, ..A, 1..A));
487            let expected: u64 = (0..n).map(|i| (a * i + b) / m).sum();
488            let result = floor_sum(n, a, b, m);
489            assert_eq!(expected, result);
490
491            let expected = (0..n).map(|i| (a * i + b) % m).min().unwrap_or(m);
492            let result = min_of_mod_of_linear(n, a, b, m);
493            assert_eq!(expected, result);
494
495            let (mut l, mut r, a, b) = rng.random((-B..B, -B..B, -B..B, -B..B));
496            if l > r {
497                swap(&mut l, &mut r);
498            }
499            let expected: i64 = (l..r).map(|i| (a * i + b).div_euclid(m as i64)).sum();
500            let result = floor_sum_i64(l, r, a, b, m);
501            assert_eq!(expected, result);
502
503            let (mut lv, mut rv) = rng.random((0..m as i64, 0..m as i64));
504            if lv > rv {
505                swap(&mut lv, &mut rv);
506            }
507            let range = lv..rv + 1;
508            let expected = (l..r)
509                .filter(|&i| range.contains(&(a * i + b).rem_euclid(m as i64)))
510                .count() as i64;
511            let result = floor_sum_range_freq(l, r, a, b, m, range);
512            assert_eq!(expected, result);
513        }
514    }
515
516    #[test]
517    fn test_floor_sum_polynomial() {
518        const P: usize = 3;
519        const A: u64 = 100;
520        const B: i64 = 100;
521        const Q: usize = 1_000;
522        let mut rng = Xorshift::default();
523        for _ in 0..Q {
524            let (n, a, b, m) = rng.random((..A, ..A, ..A, 1..A));
525            let mut expected: [[u64; P]; P] = [[0; P]; P];
526            for (x, expected) in expected.iter_mut().enumerate() {
527                for (y, expected) in expected.iter_mut().enumerate() {
528                    *expected = (0..n)
529                        .map(|i| i.pow(x as u32) * ((a * i + b) / m).pow(y as u32))
530                        .sum();
531                }
532            }
533            let result = floor_sum_polynomial::<u64, P, P>(n, a, b, m);
534            assert_eq!(expected, result);
535
536            let (mut l, mut r, a, b) = rng.random((-B..B, -B..B, -B..B, -B..B));
537            if l > r {
538                swap(&mut l, &mut r);
539            }
540            let mut expected: [[i64; P]; P] = [[0; P]; P];
541            for (x, expected) in expected.iter_mut().enumerate() {
542                for (y, expected) in expected.iter_mut().enumerate() {
543                    *expected = (l..r)
544                        .map(|i| i.pow(x as u32) * (a * i + b).div_euclid(m as i64).pow(y as u32))
545                        .sum();
546                }
547            }
548            let result = floor_sum_polynomial_i64::<i64, P, P>(l, r, a, b, m);
549            assert_eq!(expected, result);
550        }
551    }
552
553    #[test]
554    fn test_floor_power_sum() {
555        const A: u64 = 1_000;
556        const Q: usize = 1_000;
557        let mut rng = Xorshift::default();
558        for _ in 0..Q {
559            let (n, a, b, m) = rng.random((..A, ..A, ..A, 1..A));
560            let x: MInt998244353 = rng.random(..);
561            let y: MInt998244353 = rng.random(..);
562            let expected: MInt998244353 = (0..n)
563                .map(|i| {
564                    let floor = (a * i + b) / m;
565                    x.pow(i as _) * y.pow(floor as _)
566                })
567                .sum();
568            let result = floor_power_sum::<AddMulOperation<MInt998244353>>(x, y, n, a, b, m);
569            assert_eq!(expected, result);
570        }
571    }
572}
competitive/math/floor_sum.rs

competitive/math/
floor_sum.rs
