atcoder/lib/
floor_sum.rs

1//! 床関数の和に関するアルゴリズムです。
2
3/// 等差数列の各項を `m` で割った商（床）の総和を返します。
4///
5/// 次の値を `O(log m)` で計算します。
6///
7/// ```text
8/// sum_{i = 0}^{n - 1} floor((a * i + b) / m)
9/// ```
10///
11/// `n` は項数（和の対象は `i = 0..n` の半開区間）、`m` は除数、`a` は係数、`b` は定数項です。
12/// すべての引数は `u64` なので非負であり、`m` は正でなければなりません。
13///
14/// 計算量は `O(log m)` です。戻り値は `u128` ですが、非常に大きな入力では総和が `u128` に
15/// 収まらないことがあります。正しい結果を得るには総和が `u128` に収まり、計算途中の値も
16/// `u128` の範囲内であることが必要です。積の計算には `u128` を使います。
17///
18/// # Panics
19///
20/// `m` が `0` の場合は panic します。
21///
22/// # Examples
23///
24/// ```
25/// use atcoder::floor_sum::floor_sum;
26///
27/// let sum = floor_sum(6, 5, 3, 2);
28///
29/// // floor(2 / 5) + floor(5 / 5) + floor(8 / 5)
30/// // + floor(11 / 5) + floor(14 / 5) + floor(17 / 5)
31/// // = 0 + 1 + 1 + 2 + 2 + 3
32/// assert_eq!(sum, 9);
33/// ```
34pub fn floor_sum(mut n: u64, mut m: u64, mut a: u64, mut b: u64) -> u128 {
35    assert!(m > 0, "modulus must be positive");
36    if n == 0 {
37        return 0;
38    }
39
40    let mut sum = 0u128;
41
42    loop {
43        // a = q * m + r と分解する。
44        //
45        // floor((a * i + b) / m)
46        // = q * i + floor((r * i + b) / m)
47        //
48        // であるため、q * (0 + 1 + ... + n - 1) を先に加える。
49        if a >= m {
50            let quotient = a / m;
51
52            sum += n as u128 * (n - 1) as u128 / 2 * quotient as u128;
53
54            a %= m;
55        }
56
57        // b = q * m + r と分解する。
58        //
59        // 各項に q が含まれるため、q * n を先に加える。
60        if b >= m {
61            let quotient = b / m;
62
63            sum += n as u128 * quotient as u128;
64            b %= m;
65        }
66
67        // 半開区間の右端 i = n における一次式の値を求める。
68        let y_max = a as u128 * n as u128 + b as u128;
69
70        // a * n + b < m なら、残りの各項の商はすべて 0 になる。
71        if y_max < m as u128 {
72            break;
73        }
74
75        // 直線の下にある格子点を、列ごとではなく行ごとに数え直す。
76        //
77        // floor_sum(n, m, a, b)
78        // = floor_sum(
79        //     floor((a * n + b) / m),
80        //     a,
81        //     m,
82        //     (a * n + b) mod m,
83        //   )
84        //
85        // という変形を行う。
86        n = (y_max / m as u128) as u64;
87        b = (y_max % m as u128) as u64;
88
89        std::mem::swap(&mut m, &mut a);
90    }
91
92    sum
93}
94
95#[cfg(test)]
96mod tests {
97    use super::floor_sum;
98
99    /// 定義どおりに `O(n)` で床関数の和を計算します。
100    fn floor_sum_naive(n: u64, m: u64, a: u64, b: u64) -> u128 {
101        (0..n)
102            .map(|i| (a as u128 * i as u128 + b as u128) / m as u128)
103            .sum()
104    }
105
106    #[test]
107    fn calculates_floor_sum() {
108        assert_eq!(floor_sum(6, 5, 3, 2), 9);
109        assert_eq!(floor_sum(0, 5, 8, 7), 0);
110        assert_eq!(floor_sum(10, 7, 0, 15), 20);
111    }
112
113    #[test]
114    fn matches_naive_calculation() {
115        for n in 0..20 {
116            for m in 1..20 {
117                for a in 0..20 {
118                    for b in 0..20 {
119                        assert_eq!(
120                            floor_sum(n, m, a, b),
121                            floor_sum_naive(n, m, a, b),
122                            "n = {n}, m = {m}, a = {a}, b = {b}",
123                        );
124                    }
125                }
126            }
127        }
128    }
129}
atcoder/lib/floor_sum.rs

atcoder/lib/
floor_sum.rs
