Yosupo Library Checker - Sum of Floor of Linear

floor sum の verify！

• 問題へのリンク

問題概要

正の整数 が与えられる。

の値を求めよ ( ケース与えられる)。

制約

• 
• 
• 

解法

ACL practice C 問題の解法、および、ACL の実装を参考にして実装した。

atcoder.jp

コード

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

// sum_{i=0}^{n-1} floor((a * i + b) / m)
template<class T> T floor_sum(T n, T a, T b, T m) {
    assert(n >= 0 && m >= 1);
    T res = 0;
    if (a < 0) {
        T a2 = (a % m + m) % m;
        res -= n * (n - 1) / 2 * ((a2 - a) / m);
        a = a2;
    }
    if (b < 0) {
        T b2 = (b % m + m) % m;
        res -= n * ((b2 - b) / m);
        b = b2;
    }
    
    while (true) {
        if (a >= m) {
            res += n * (n - 1) / 2 * (a / m);
            a %= m;
        }
        if (b >= m) {
            res += n * (b / m);
            b %= m;
        }
        T y_max = a * n + b;
        if (y_max < m) break;
        n = y_max / m;
        b = y_max % m;
        swap(m, a);
    }
    return res;
}

// Library Checker
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;
    }
}