けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder Library Practice Contest C - Floor Sum (1D)

Floor Sum が AtCoder 標準ライブラリに入ったのはびっくりした。結構マニアックな印象だった。

問題概要

この問題は  T ケース与えられる。整数  N,M,A,B が与えられるので、

 \displaystyle \sum_{i=0}^{N-1} \mathrm{floor}\bigl(\frac{Ai + B}{M} \bigr)

を求めよ。

制約

  •  1 \le T \le 10^{5}
  •  1 \le N, M \le 10^{9}
  •  0 \le A, B \lt M

考えたこと

等差数列を  M で割った商の総和を求めたくなることは、確かに時々ある。具体的には

 \displaystyle \sum_{i=0}^{N-1} \mathrm{floor}\bigl(\frac{Ai + B}{M} \bigr)

を求めたい。とても難しいタスクだけど、 O(N + M + A + B) の計算量でできてしまう!!

なお floor sum については公式解説がある!

atcoder.jp

また、この公式解説の元ネタとなった記事もある。

rsk0315.hatenablog.com

他にも、kyopro_friends さんのツイートもわかりやすい!

コード

#include <bits/stdc++.h>
#include <atcoder/math>
using namespace std;

int main() {
    int T;
    cin >> T;
    while (T--) {
        long long N, M, A, B;
        cin >> N >> M >> A >> B;
        cout << atcoder::floor_sum(N, M, A, B) << endl;
    }
}

自前ライブラリでも AC

#include <bits/stdc++.h>
#include <atcoder/math>
using namespace std;

// sum_{i=0}^{n-1} floor((a * i + b) / m)
// O(log(n + m + a + b))
long long floor_sum(long long n, long long a, long long b, long long m) {
    if (n == 0) return 0;
    long long res = 0;
    if (a >= m) {
        res += n * (n - 1) * (a / m) / 2;
        a %= m;
    }
    if (b >= m) {
        res += n * (b / m);
        b %= m;
    }
    if (a == 0) return res;
    long long ymax = (a * n + b) / m, xmax = ymax * m - b;
    if (ymax == 0) return res;
    res += (n - (xmax + a - 1) / a) * ymax;
    res += floor_sum(ymax, m, (a - xmax % a) % a, a);
    return res;
}

int main() {
    int T;
    cin >> T;
    while (T--) {
        long long N, M, A, B;
        cin >> N >> M >> A >> B;
        cout << floor_sum(N, A, B, M) << endl;
    }
}