nekolib/math/
linear_floor_sum.rs

1//! $ \\sum\_{i=0}\^{n-1} \\left\\lfloor\\frac{ai+b}{m}\\right\\rfloor. $
2
3/// $ \\sum\_{i=0}\^{n-1} \\left\\lfloor\\frac{ai+b}{m}\\right\\rfloor. $
4///
5/// # Requirements
6/// - $n \\ge 0$
7/// - $m \\gt 0$
8///
9/// # Idea
10/// あとで書く。
11///
12/// ## See also
13/// - <https://rsk0315.hatenablog.com/entry/2020/12/13/231307>
14/// - <https://atcoder.jp/contests/practice2/editorial/579>
15///
16/// # Examples
17/// ```
18/// use nekolib::math::LinearFloorSum;
19///
20/// assert_eq!(4_u128.linear_floor_sum(10, 6, 3), 3);
21/// assert_eq!(6_u128.linear_floor_sum(5, 4, 3), 13);
22/// assert_eq!(1_u128.linear_floor_sum(1, 0, 0), 0);
23/// assert_eq!(31415_u128.linear_floor_sum(92653, 58979, 32384), 314095480);
24/// assert_eq!(
25///     1000000000_u128.linear_floor_sum(1000000000, 999999999, 999999999),
26///     499999999500000000
27/// );
28/// assert_eq!(14_i128.linear_floor_sum(23, -7, -39), -58);
29/// ```
30pub trait LinearFloorSum {
31    fn linear_floor_sum(self, m: Self, a: Self, b: Self) -> Self;
32}
33
34macro_rules! impl_uint {
35    ($t:ty) => {
36        impl LinearFloorSum for $t {
37            fn linear_floor_sum(
38                self,
39                mut m: Self,
40                mut a: Self,
41                mut b: Self
42            ) -> Self {
43                let mut n = self;
44                let mut res = 0;
45                loop {
46                    if a >= m {
47                        res += n * (n - 1) / 2 * (a / m);
48                        a %= m;
49                    }
50                    if b >= m {
51                        res += n * (b / m);
52                        b %= m;
53                    }
54                    let y = a * n + b;
55                    if y < m {
56                        break;
57                    }
58                    n = y / m;
59                    b = y % m;
60                    std::mem::swap(&mut m, &mut a);
61                }
62                res
63            }
64        }
65    };
66    ( $($t:ty)* ) => { $(impl_uint!($t);)* };
67}
68
69macro_rules! impl_int {
70    ($t:ty) => {
71        impl LinearFloorSum for $t {
72            fn linear_floor_sum(self, m: Self, a: Self, b: Self) -> Self {
73                let n = self;
74                let mut res = 0;
75                assert!(m > 0);
76                assert!(n >= 0);
77                let a = if a < 0 {
78                    let a2 = a.rem_euclid(m);
79                    res -= n * (n - 1) / 2 * ((a2 - a) / m);
80                    a2 as u128
81                } else {
82                    a as u128
83                };
84                let b = if b < 0 {
85                    let b2 = b.rem_euclid(m);
86                    res -= n * ((b2 - b) / m);
87                    b2 as u128
88                } else {
89                    b as u128
90                };
91                res + (self as u128).linear_floor_sum(m as u128, a, b) as $t
92            }
93        }
94    };
95    ( $($t:ty)* ) => { $(impl_int!($t);)* };
96}
97
98impl_uint!(u8 u16 u32 u64 u128 usize);
99impl_int!(i8 i16 i32 i64 i128 isize);
nekolib/math/linear_floor_sum.rs

nekolib/math/
linear_floor_sum.rs
