AtCoder Library Practice Contest C - Floor Sum

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

• 問題へのリンク

問題概要

この問題は ケース与えられる。整数 が与えられるので、

を求めよ。

制約

• 
• 
• 

考えたこと

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

を求めたい。とても難しいタスクだけど、 の計算量でできてしまう！！

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

atcoder.jp

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

rsk0315.hatenablog.com

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

サーバル「ALPC C『floor sum』を解いたよ。問題は、「直線 y=(Ax+B)/Mの下にある格子点の個数を求めよ」になるよ。境界を含むか含まないかには注意してね。後で話を簡単にするために左右を反転して考えることにしておくよ」 pic.twitter.com/fZvOZyOWPu

— 競技プログラミングをするフレンズ (@kyopro_friends) 2020年9月10日

コード

#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;
    }
}