competitive_programming_rs/math/
floor_sum.rs

1pub fn floor_sum(n: i64, m: i64, mut a: i64, mut b: i64) -> i64 {
2    let mut ans = 0;
3    if a >= m {
4        ans += (n - 1) * n * (a / m) / 2;
5        a %= m;
6    }
7    if b >= m {
8        ans += n * (b / m);
9        b %= m;
10    }
11
12    let y_max = (a * n + b) / m;
13    let x_max = y_max * m - b;
14    if y_max == 0 {
15        ans
16    } else {
17        ans += (n - (x_max + a - 1) / a) * y_max;
18        ans += floor_sum(y_max, a, m, (a - x_max % a) % a);
19        ans
20    }
21}
22
23#[cfg(test)]
24mod tests {
25    use super::*;
26
27    #[test]
28    fn test_floor_sum() {
29        assert_eq!(3, floor_sum(4, 10, 6, 3));
30        assert_eq!(13, floor_sum(6, 5, 4, 3));
31        assert_eq!(0, floor_sum(1, 1, 0, 0));
32        assert_eq!(314095480, floor_sum(31415, 92653, 58979, 32384));
33        assert_eq!(
34            499999999500000000,
35            floor_sum(1000000000, 1000000000, 999999999, 999999999)
36        );
37    }
38}