verify/yosupo-sum-of-floor-of-linear.test.cpp

• View this file on GitHub
• Last update: 2026-02-07 06:46:25+00:00
• Problem: https://judge.yosupo.jp/problem/sum_of_floor_of_linear

Depends on

• floor sum (math/number-theory/floor-sum.hpp)

Code

// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/sum_of_floor_of_linear

#include <cassert>
#include <cstdint>
#include <iostream>

#include "../math/number-theory/floor-sum.hpp"

#ifdef __SIZEOF_INT128__
static __int128 floor_div_i128(__int128 x, __int128 m) {
    assert(m > 0);
    __int128 q = x / m;
    __int128 r = x % m;
    if (r < 0)
        --q;
    return q;
}

static __int128 brute_floor_sum_i128(long long n, long long m, long long a,
                                     long long b) {
    assert(n >= 0);
    assert(m > 0);
    __int128 res = 0;
    for (long long i = 0; i < n; ++i) {
        res += floor_div_i128(static_cast<__int128>(a) * i + b, m);
    }
    return res;
}

static std::uint64_t xorshift64(std::uint64_t &x) {
    x ^= x << 7;
    x ^= x >> 9;
    return x;
}

static void self_check_negative() {
    std::uint64_t rng = 88172645463325252ull;

    for (int iter = 0; iter < 5000; ++iter) {
        const long long n = static_cast<long long>(xorshift64(rng) % 25);
        const long long m = static_cast<long long>(xorshift64(rng) % 25) + 1;
        const long long a = static_cast<long long>(xorshift64(rng) % 81) - 40;
        const long long b = static_cast<long long>(xorshift64(rng) % 81) - 40;

        const __int128 want = brute_floor_sum_i128(n, m, a, b);
        const __int128 got1 = floor_sum<__int128>(
            static_cast<__int128>(n), static_cast<__int128>(m),
            static_cast<__int128>(a), static_cast<__int128>(b));
        assert(want == got1);

        const long long got2 = floor_sum<long long>(n, m, a, b);
        assert(want == static_cast<__int128>(got2));
    }

    {
        const long long n = 0, m = 7, a = -3, b = -5;
        assert(floor_sum<long long>(n, m, a, b) == 0);
    }
    {
        const long long n = 10, m = 1, a = -7, b = 3;
        const __int128 want = brute_floor_sum_i128(n, m, a, b);
        const long long got = floor_sum<long long>(n, m, a, b);
        assert(want == static_cast<__int128>(got));
    }
}
#endif

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

#ifdef __SIZEOF_INT128__
    self_check_negative();
#endif

    int T;
    std::cin >> T;
    while (T--) {
        long long n, m, a, b;
        std::cin >> n >> m >> a >> b;
        std::cout << floor_sum<long long>(n, m, a, b) << '\n';
    }
    return 0;
}

#line 1 "verify/yosupo-sum-of-floor-of-linear.test.cpp"
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/sum_of_floor_of_linear

#include <cassert>
#include <cstdint>
#include <iostream>

#line 1 "math/number-theory/floor-sum.hpp"



// 一次式の床関数の和 Σ_{i=0}^{n-1} floor((a i + b) / m) を求める。
// n は非負、m は正を仮定する。
// a, b は負でもよい（数学的な床除算で扱う）。
// T は整数型（__int128 を含む）。
// n, m が sqrt(max(T)) 程度なら a*n+b などの中間計算が安全になりやすい。
// 計算量 O(log m)。

#line 12 "math/number-theory/floor-sum.hpp"
#include <type_traits>
#include <utility>

namespace floor_sum_internal {

template <class T> struct is_integral : std::is_integral<T> {};

#ifdef __SIZEOF_INT128__
template <> struct is_integral<__int128_t> : std::true_type {};
template <> struct is_integral<__uint128_t> : std::true_type {};
#endif

template <class T> struct is_signed : std::is_signed<T> {};

#ifdef __SIZEOF_INT128__
template <> struct is_signed<__int128_t> : std::true_type {};
template <> struct is_signed<__uint128_t> : std::false_type {};
#endif

template <class T> T floor_div(T x, T y) {
    assert(y > 0);
    if constexpr (is_signed<T>::value) {
        T q = x / y;
        T r = x % y;
        if (r < 0)
            --q;
        return q;
    } else {
        return x / y;
    }
}

template <class T> T floor_mod(T x, T y) {
    assert(y > 0);
    if constexpr (is_signed<T>::value) {
        T r = x % y;
        if (r < 0)
            r += y;
        return r;
    } else {
        return x % y;
    }
}

// Σ_{i=0}^{n-1} i = n(n-1)/2
template <class T> T sum_0_to_n_minus_1(T n) {
    if (n == 0)
        return 0;
    if ((n & 1) == 0)
        return (n / 2) * (n - 1);
    return n * ((n - 1) / 2);
}

} // namespace floor_sum_internal

template <class T> T floor_sum(T n, T m, T a, T b) {
    static_assert(floor_sum_internal::is_integral<T>::value,
                  "T must be integer.");
    if constexpr (floor_sum_internal::is_signed<T>::value)
        assert(n >= 0);
    assert(m > 0);
    if (n == 0)
        return 0;

    T ans = 0;

    {
        const T q = floor_sum_internal::floor_div(a, m);
        a = floor_sum_internal::floor_mod(a, m);
        ans += q * floor_sum_internal::sum_0_to_n_minus_1(n);
    }
    {
        const T q = floor_sum_internal::floor_div(b, m);
        b = floor_sum_internal::floor_mod(b, m);
        ans += q * n;
    }

    while (true) {
        if (a >= m) {
            const T q = a / m;
            ans += q * floor_sum_internal::sum_0_to_n_minus_1(n);
            a %= m;
        }
        if (b >= m) {
            const T q = b / m;
            ans += q * n;
            b %= m;
        }

        const T y_max = a * n + b;
        if (y_max < m)
            break;

        n = y_max / m;
        b = y_max % m;
        std::swap(m, a);
    }
    return ans;
}


#line 8 "verify/yosupo-sum-of-floor-of-linear.test.cpp"

#ifdef __SIZEOF_INT128__
static __int128 floor_div_i128(__int128 x, __int128 m) {
    assert(m > 0);
    __int128 q = x / m;
    __int128 r = x % m;
    if (r < 0)
        --q;
    return q;
}

static __int128 brute_floor_sum_i128(long long n, long long m, long long a,
                                     long long b) {
    assert(n >= 0);
    assert(m > 0);
    __int128 res = 0;
    for (long long i = 0; i < n; ++i) {
        res += floor_div_i128(static_cast<__int128>(a) * i + b, m);
    }
    return res;
}

static std::uint64_t xorshift64(std::uint64_t &x) {
    x ^= x << 7;
    x ^= x >> 9;
    return x;
}

static void self_check_negative() {
    std::uint64_t rng = 88172645463325252ull;

    for (int iter = 0; iter < 5000; ++iter) {
        const long long n = static_cast<long long>(xorshift64(rng) % 25);
        const long long m = static_cast<long long>(xorshift64(rng) % 25) + 1;
        const long long a = static_cast<long long>(xorshift64(rng) % 81) - 40;
        const long long b = static_cast<long long>(xorshift64(rng) % 81) - 40;

        const __int128 want = brute_floor_sum_i128(n, m, a, b);
        const __int128 got1 = floor_sum<__int128>(
            static_cast<__int128>(n), static_cast<__int128>(m),
            static_cast<__int128>(a), static_cast<__int128>(b));
        assert(want == got1);

        const long long got2 = floor_sum<long long>(n, m, a, b);
        assert(want == static_cast<__int128>(got2));
    }

    {
        const long long n = 0, m = 7, a = -3, b = -5;
        assert(floor_sum<long long>(n, m, a, b) == 0);
    }
    {
        const long long n = 10, m = 1, a = -7, b = 3;
        const __int128 want = brute_floor_sum_i128(n, m, a, b);
        const long long got = floor_sum<long long>(n, m, a, b);
        assert(want == static_cast<__int128>(got));
    }
}
#endif

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

#ifdef __SIZEOF_INT128__
    self_check_negative();
#endif

    int T;
    std::cin >> T;
    while (T--) {
        long long n, m, a, b;
        std::cin >> n >> m >> a >> b;
        std::cout << floor_sum<long long>(n, m, a, b) << '\n';
    }
    return 0;
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	example_00	AC	7 ms	4 MB
g++	random_00	AC	19 ms	4 MB
g++	random_01	AC	58 ms	4 MB
g++	random_02	AC	46 ms	4 MB
g++	random_03	AC	34 ms	4 MB
g++	random_04	AC	18 ms	4 MB
g++	small_00	AC	13 ms	4 MB
g++	small_01	AC	32 ms	4 MB
g++	small_02	AC	27 ms	4 MB
g++	small_03	AC	20 ms	4 MB
g++	small_04	AC	12 ms	4 MB
clang++	example_00	AC	7 ms	4 MB
clang++	random_00	AC	20 ms	4 MB
clang++	random_01	AC	59 ms	4 MB
clang++	random_02	AC	47 ms	4 MB
clang++	random_03	AC	34 ms	4 MB
clang++	random_04	AC	18 ms	4 MB
clang++	small_00	AC	13 ms	4 MB
clang++	small_01	AC	33 ms	4 MB
clang++	small_02	AC	27 ms	4 MB
clang++	small_03	AC	20 ms	4 MB
clang++	small_04	AC	12 ms	4 MB

Back to top page