haar_lib/math/
sum_floor_linear.rs

1//! $\sum_{i=0}^{n-1} \lfloor \frac{ai+b}{m} \rfloor$
2
3/// $\sum_{i=0}^{n-1} \lfloor \frac{ai+b}{m} \rfloor$
4pub fn sum_floor_linear(n: u64, m: u64, mut a: u64, mut b: u64) -> u64 {
5    let mut ret = 0;
6
7    if b >= m {
8        ret += n * (b / m);
9        b %= m;
10    }
11
12    if a >= m {
13        ret += n * (n - 1) * (a / m) / 2;
14        a %= m;
15    }
16
17    let y_max = (a * n + b) / m;
18    if y_max == 0 {
19        return ret;
20    }
21
22    let x_max = y_max * m - b;
23
24    ret += (n - x_max.div_ceil(a)) * y_max;
25    ret += sum_floor_linear(y_max, a, m, (a - x_max % a) % a);
26
27    ret
28}
29
30#[cfg(test)]
31mod tests {
32    use super::*;
33    #[test]
34    fn test() {
35        // https://judge.yosupo.jp/problem/sum_of_floor_of_linear
36        assert_eq!(sum_floor_linear(4, 10, 6, 3), 3);
37        assert_eq!(sum_floor_linear(6, 5, 4, 3), 13);
38        assert_eq!(sum_floor_linear(1, 1, 0, 0), 0);
39        assert_eq!(sum_floor_linear(31415, 92653, 58979, 32384), 314095480);
40        assert_eq!(
41            sum_floor_linear(1000000000, 1000000000, 999999999, 999999999),
42            499999999500000000
43        );
44    }
45}
haar_lib/math/sum_floor_linear.rs

haar_lib/math/
sum_floor_linear.rs
