verify/standalone-generalized-floor-sum-degree-le-2.test.cpp

• View this file on GitHub
• Last update: 2026-07-11 16:54:42+09:00

Depends on

• 128 bit 整数型 (internal/int128.hpp)
• 一般化 floor sum（次数 2 以下） (math/number-theory/generalized-floor-sum-degree-le-2.hpp)

Code

// competitive-verifier: STANDALONE

#include <cassert>
#include <cstdint>

#include "../math/number-theory/generalized-floor-sum-degree-le-2.hpp"

long long floor_div_brute(long long x, long long y) {
    assert(y > 0);

    long long q = x / y;
    long long r = x % y;
    if (r < 0) {
        --q;
    }
    return q;
}

int main() {
    for (long long n = 0; n <= 20; ++n) {
        for (long long m = 1; m <= 20; ++m) {
            for (long long a = -20; a <= 20; ++a) {
                for (long long b = -20; b <= 20; ++b) {
                    const auto res =
                        generalized_floor_sum_degree_le_2<long long>(n, m, a,
                                                                     b);

                    long long s_01 = 0, s_11 = 0, s_02 = 0;
                    for (long long i = 0; i < n; ++i) {
                        const long long value = floor_div_brute(a * i + b, m);
                        s_01 += value;
                        s_11 += i * value;
                        s_02 += value * value;
                    }

                    assert(res.ans_01 == s_01);
                    assert(res.ans_11 == s_11);
                    assert(res.ans_02 == s_02);
                }
            }
        }
    }

    {
        using T = std::uint64_t;

        const T n = 900931385;
        const T m = 333006410;
        const T x = 263208878;
        const T y = 243209245;

        const auto r0 = generalized_floor_sum_degree_le_2<T>(n, m, x, T(0));
        const auto r1 = generalized_floor_sum_degree_le_2<T>(n, m, x, y);

        T ans = 0;
        ans -= r0.ans_01 * n - r0.ans_11;
        ans -= r1.ans_01 * (n - 1) - 2 * r1.ans_11;

        assert(ans == 202919340569980512ULL);
    }

    return 0;
}

#line 1 "verify/standalone-generalized-floor-sum-degree-le-2.test.cpp"
// competitive-verifier: STANDALONE

#include <cassert>
#include <cstdint>

#line 1 "math/number-theory/generalized-floor-sum-degree-le-2.hpp"



// 一般化 floor sum（次数 2 以下）のうち、（0, 1）、（1, 1）、（0, 2）をまとめて求める。
// ans_01 = Σ floor((a i + b) / m)、ans_11 = Σ i floor((a i + b) / m)、
// ans_02 = Σ floor((a i + b) / m)^2。
// n >= 0、m > 0 を仮定する。T が符号付きの場合、a, b は負でもよい。
// T および明示的に指定する Internal は整数型であり、内部計算が内部型の範囲に収まることを仮定する。
// 計算量 O(log m)。

#line 12 "math/number-theory/generalized-floor-sum-degree-le-2.hpp"
#include <limits>
#include <type_traits>
#include <utility>

#line 1 "internal/int128.hpp"



// 128 bit 符号なし整数型 UInt128 と 128 bit 符号付き整数型 Int128 を提供する。
// GCC 拡張には依存しない。

#include <bit>
#line 10 "internal/int128.hpp"
#include <istream>
#line 12 "internal/int128.hpp"
#include <ostream>
#include <string>
#line 15 "internal/int128.hpp"

namespace NicheLibrary {
class UInt128 {
  public:
    constexpr UInt128() = default;

    template <
        class T,
        std::enable_if_t<std::is_integral_v<T> && std::is_signed_v<T>, int> = 0>
    constexpr UInt128(T value)
        : high_(value < 0 ? ~std::uint64_t{} : std::uint64_t{}),
          low_(static_cast<std::uint64_t>(value)) {}

    template <class T,
              std::enable_if_t<std::is_integral_v<T> && !std::is_signed_v<T>,
                               int> = 0>
    constexpr UInt128(T value)
        : high_(0), low_(static_cast<std::uint64_t>(value)) {}

    static constexpr UInt128 from_words(std::uint64_t high, std::uint64_t low) {
        UInt128 value;
        value.high_ = high;
        value.low_ = low;
        return value;
    }

    constexpr std::uint64_t high() const { return high_; }
    constexpr std::uint64_t low() const { return low_; }

    template <class T, std::enable_if_t<std::is_integral_v<T>, int> = 0>
    explicit constexpr operator T() const {
        return static_cast<T>(low_);
    }

    explicit constexpr operator bool() const { return high_ != 0 || low_ != 0; }

    explicit constexpr operator long double() const {
        if (high_ == 0) {
            return static_cast<long double>(low_);
        }
        return static_cast<long double>(high_) * 0x1p64L +
               static_cast<long double>(low_);
    }

    friend constexpr bool operator==(UInt128 lhs, UInt128 rhs) {
        return lhs.high_ == rhs.high_ && lhs.low_ == rhs.low_;
    }

    friend constexpr bool operator!=(UInt128 lhs, UInt128 rhs) {
        return !(lhs == rhs);
    }

    friend constexpr bool operator<(UInt128 lhs, UInt128 rhs) {
        if (lhs.high_ != rhs.high_) {
            return lhs.high_ < rhs.high_;
        }
        return lhs.low_ < rhs.low_;
    }

    friend constexpr bool operator>(UInt128 lhs, UInt128 rhs) {
        return rhs < lhs;
    }

    friend constexpr bool operator<=(UInt128 lhs, UInt128 rhs) {
        return !(rhs < lhs);
    }

    friend constexpr bool operator>=(UInt128 lhs, UInt128 rhs) {
        return !(lhs < rhs);
    }

    constexpr UInt128 operator+() const { return *this; }

    constexpr UInt128 operator-() const { return UInt128{} - *this; }

    constexpr UInt128 &operator+=(UInt128 rhs) {
        const std::uint64_t old_low = low_;
        low_ += rhs.low_;
        high_ += rhs.high_ + (low_ < old_low ? 1 : 0);
        return *this;
    }

    constexpr UInt128 &operator-=(UInt128 rhs) {
        const std::uint64_t old_low = low_;
        low_ -= rhs.low_;
        high_ -= rhs.high_ + (old_low < rhs.low_ ? 1 : 0);
        return *this;
    }

    constexpr UInt128 &operator*=(UInt128 rhs) {
        *this = *this * rhs;
        return *this;
    }

    constexpr UInt128 &operator/=(UInt128 rhs) {
        *this = *this / rhs;
        return *this;
    }

    constexpr UInt128 &operator%=(UInt128 rhs) {
        *this = *this % rhs;
        return *this;
    }

    friend constexpr UInt128 operator+(UInt128 lhs, UInt128 rhs) {
        lhs += rhs;
        return lhs;
    }

    friend constexpr UInt128 operator-(UInt128 lhs, UInt128 rhs) {
        lhs -= rhs;
        return lhs;
    }

    friend constexpr UInt128 operator*(UInt128 lhs, UInt128 rhs) {
        const UInt128 p00 = multiply_u64(lhs.low_, rhs.low_);
        return from_words(
            p00.high_ + lhs.low_ * rhs.high_ + lhs.high_ * rhs.low_, p00.low_);
    }

    static constexpr void div_mod(UInt128 lhs, UInt128 rhs, UInt128 &quotient,
                                  UInt128 &remainder) {
        assert(rhs != UInt128{});
        div_mod_unchecked(lhs, rhs, quotient, remainder);
    }

    friend constexpr UInt128 operator/(UInt128 lhs, UInt128 rhs) {
        UInt128 quotient;
        UInt128 remainder;
        div_mod(lhs, rhs, quotient, remainder);
        return quotient;
    }

    friend constexpr UInt128 operator%(UInt128 lhs, UInt128 rhs) {
        UInt128 quotient;
        UInt128 remainder;
        div_mod(lhs, rhs, quotient, remainder);
        return remainder;
    }

  private:
    friend class Int128;

    static constexpr std::uint64_t word_base() {
        return std::uint64_t{1} << 32;
    }

    static constexpr std::uint64_t word_mask() { return word_base() - 1; }

    static constexpr void to_words(UInt128 value, std::uint32_t words[4]) {
        words[0] = static_cast<std::uint32_t>(value.low_ & word_mask());
        words[1] = static_cast<std::uint32_t>(value.low_ >> 32);
        words[2] = static_cast<std::uint32_t>(value.high_ & word_mask());
        words[3] = static_cast<std::uint32_t>(value.high_ >> 32);
    }

    static constexpr UInt128 from_32bit_words(const std::uint32_t words[4]) {
        return UInt128::from_words(
            (static_cast<std::uint64_t>(words[3]) << 32) | words[2],
            (static_cast<std::uint64_t>(words[1]) << 32) | words[0]);
    }

    static constexpr int word_length(const std::uint32_t words[4]) {
        int length = 4;
        while (length > 0 && words[length - 1] == 0) {
            --length;
        }
        return length;
    }

    // divisor ≠ 0 が呼び出し元で保証されていることを仮定する。
    static constexpr UInt128 div_mod_uint32(UInt128 value,
                                            std::uint32_t divisor,
                                            std::uint32_t &remainder) {
        const std::uint32_t words[4] = {
            static_cast<std::uint32_t>(value.low_ & word_mask()),
            static_cast<std::uint32_t>(value.low_ >> 32),
            static_cast<std::uint32_t>(value.high_ & word_mask()),
            static_cast<std::uint32_t>(value.high_ >> 32),
        };

        std::uint32_t quotient_words[4] = {};
        std::uint64_t rem = 0;
        for (int i = 3; i >= 0; --i) {
            const std::uint64_t current = (rem << 32) | words[i];
            quotient_words[i] = static_cast<std::uint32_t>(current / divisor);
            rem = current % divisor;
        }

        remainder = static_cast<std::uint32_t>(rem);
        return from_32bit_words(quotient_words);
    }

    static constexpr void shift_left_words(const std::uint32_t source[4],
                                           int source_length, int shift,
                                           std::uint32_t result[5]) {
        std::uint64_t carry = 0;

        for (int i = 0; i < source_length; ++i) {
            const std::uint64_t current =
                (static_cast<std::uint64_t>(source[i]) << shift) | carry;
            result[i] = static_cast<std::uint32_t>(current & word_mask());
            carry = current >> 32;
        }

        result[source_length] = static_cast<std::uint32_t>(carry);
    }

    static constexpr void shift_right_words(const std::uint32_t source[5],
                                            int source_length, int shift,
                                            std::uint32_t result[4]) {
        if (shift == 0) {
            for (int i = 0; i < source_length; ++i) {
                result[i] = source[i];
            }
            return;
        }

        for (int i = 0; i < source_length; ++i) {
            const std::uint32_t next = i + 1 < 5 ? source[i + 1] : 0;
            result[i] = static_cast<std::uint32_t>(
                (source[i] >> shift) |
                (static_cast<std::uint64_t>(next) << (32 - shift)));
        }
    }

    static constexpr bool subtract_product(std::uint32_t words[5], int offset,
                                           const std::uint32_t divisor[5],
                                           int divisor_length,
                                           std::uint32_t multiplier) {
        std::uint64_t borrow = 0;
        std::uint64_t carry = 0;

        for (int i = 0; i < divisor_length; ++i) {
            const std::uint64_t product =
                static_cast<std::uint64_t>(multiplier) * divisor[i] + carry;
            carry = product >> 32;
            const std::uint64_t subtrahend = (product & word_mask()) + borrow;
            if (static_cast<std::uint64_t>(words[offset + i]) >= subtrahend) {
                words[offset + i] = static_cast<std::uint32_t>(
                    static_cast<std::uint64_t>(words[offset + i]) - subtrahend);
                borrow = 0;
            } else {
                words[offset + i] = static_cast<std::uint32_t>(
                    word_base() + words[offset + i] - subtrahend);
                borrow = 1;
            }
        }

        const std::uint64_t subtrahend = carry + borrow;
        if (static_cast<std::uint64_t>(words[offset + divisor_length]) >=
            subtrahend) {
            words[offset + divisor_length] = static_cast<std::uint32_t>(
                static_cast<std::uint64_t>(words[offset + divisor_length]) -
                subtrahend);
            return false;
        }

        words[offset + divisor_length] = static_cast<std::uint32_t>(
            word_base() + words[offset + divisor_length] - subtrahend);
        return true;
    }

    static constexpr void add_back(std::uint32_t words[5], int offset,
                                   const std::uint32_t divisor[5],
                                   int divisor_length) {
        std::uint64_t carry = 0;
        for (int i = 0; i < divisor_length; ++i) {
            const std::uint64_t sum =
                static_cast<std::uint64_t>(words[offset + i]) + divisor[i] +
                carry;
            words[offset + i] = static_cast<std::uint32_t>(sum & word_mask());
            carry = sum >> 32;
        }

        words[offset + divisor_length] = static_cast<std::uint32_t>(
            static_cast<std::uint64_t>(words[offset + divisor_length]) + carry);
    }

    static constexpr void div_mod_knuth(const std::uint32_t lhs_words[4],
                                        int lhs_length,
                                        const std::uint32_t rhs_words[4],
                                        int rhs_length,
                                        std::uint32_t quotient_words[4],
                                        std::uint32_t remainder_words[4]) {
        const int shift = std::countl_zero(rhs_words[rhs_length - 1]);
        std::uint32_t normalized_lhs[5] = {};
        std::uint32_t normalized_rhs[5] = {};
        shift_left_words(lhs_words, lhs_length, shift, normalized_lhs);
        shift_left_words(rhs_words, rhs_length, shift, normalized_rhs);

        const int quotient_length = lhs_length - rhs_length + 1;
        for (int j = quotient_length - 1; j >= 0; --j) {
            const std::uint64_t numerator =
                (static_cast<std::uint64_t>(normalized_lhs[j + rhs_length])
                 << 32) |
                normalized_lhs[j + rhs_length - 1];
            std::uint64_t estimate = numerator / normalized_rhs[rhs_length - 1];
            std::uint64_t estimate_remainder =
                numerator % normalized_rhs[rhs_length - 1];
            while (estimate == word_base() ||
                   estimate * normalized_rhs[rhs_length - 2] >
                       word_base() * estimate_remainder +
                           normalized_lhs[j + rhs_length - 2]) {
                --estimate;
                estimate_remainder += normalized_rhs[rhs_length - 1];
                if (estimate_remainder >= word_base()) {
                    break;
                }
            }
            quotient_words[j] = static_cast<std::uint32_t>(estimate);
            if (subtract_product(normalized_lhs, j, normalized_rhs, rhs_length,
                                 quotient_words[j])) {
                --quotient_words[j];
                add_back(normalized_lhs, j, normalized_rhs, rhs_length);
            }
        }

        std::uint32_t shifted_remainder[4] = {};
        shift_right_words(normalized_lhs, rhs_length, shift, shifted_remainder);
        for (int i = 0; i < rhs_length; ++i) {
            remainder_words[i] = shifted_remainder[i];
        }
    }

    // 呼び出し元で lhs >= rhs > 0 が保証されていることを仮定する。
    static constexpr void div_mod_32bit_words(UInt128 lhs, UInt128 rhs,
                                              UInt128 &quotient,
                                              UInt128 &remainder) {
        std::uint32_t lhs_words[4] = {};
        std::uint32_t rhs_words[4] = {};
        to_words(lhs, lhs_words);
        to_words(rhs, rhs_words);

        const int lhs_length = word_length(lhs_words);
        const int rhs_length = word_length(rhs_words);

        if (rhs_length == 1) {
            std::uint32_t rem = 0;
            quotient = div_mod_uint32(lhs, rhs_words[0], rem);
            remainder = UInt128(rem);
            return;
        }

        std::uint32_t quotient_words[4] = {};
        std::uint32_t remainder_words[4] = {};
        div_mod_knuth(lhs_words, lhs_length, rhs_words, rhs_length,
                      quotient_words, remainder_words);

        quotient = from_32bit_words(quotient_words);
        remainder = from_32bit_words(remainder_words);
    }

    // rhs ≠ 0 が呼び出し元で保証されていることを仮定する。
    static constexpr void div_mod_unchecked(UInt128 lhs, UInt128 rhs,
                                            UInt128 &quotient,
                                            UInt128 &remainder) {
        UInt128 q, r;
        if (rhs.high_ == 0) {
            if (lhs.high_ == 0) {
                q = UInt128(lhs.low_ / rhs.low_);
                r = UInt128(lhs.low_ % rhs.low_);
            } else if (rhs.low_ == 1) {
                q = lhs;
                r = UInt128{};
            } else if (rhs.low_ <= std::numeric_limits<std::uint32_t>::max()) {
                std::uint32_t rem = 0;
                q = div_mod_uint32(lhs, static_cast<std::uint32_t>(rhs.low_),
                                   rem);
                r = UInt128(rem);
            } else {
                div_mod_32bit_words(lhs, rhs, q, r);
            }
        } else if (lhs < rhs) {
            q = UInt128{};
            r = lhs;
        } else {
            div_mod_32bit_words(lhs, rhs, q, r);
        }
        quotient = q;
        remainder = r;
    }

    static constexpr UInt128 multiply_u64(std::uint64_t lhs,
                                          std::uint64_t rhs) {
        constexpr std::uint64_t mask = (std::uint64_t{1} << 32) - 1;
        const std::uint64_t lhs_low = lhs & mask;
        const std::uint64_t lhs_high = lhs >> 32;
        const std::uint64_t rhs_low = rhs & mask;
        const std::uint64_t rhs_high = rhs >> 32;

        const std::uint64_t p00 = lhs_low * rhs_low;
        const std::uint64_t p01 = lhs_low * rhs_high;
        const std::uint64_t p10 = lhs_high * rhs_low;
        const std::uint64_t p11 = lhs_high * rhs_high;

        const std::uint64_t middle = (p00 >> 32) + (p01 & mask) + (p10 & mask);
        return from_words(p11 + (p01 >> 32) + (p10 >> 32) + (middle >> 32),
                          (middle << 32) | (p00 & mask));
    }

    std::uint64_t high_ = 0;
    std::uint64_t low_ = 0;
};

class Int128 {
  public:
    constexpr Int128() = default;

    template <class T, std::enable_if_t<std::is_integral_v<T>, int> = 0>
    constexpr Int128(T value) : value_(value) {}

    static constexpr Int128 from_words(std::uint64_t high, std::uint64_t low) {
        Int128 value;
        value.value_ = UInt128::from_words(high, low);
        return value;
    }

    constexpr std::uint64_t high() const { return value_.high(); }
    constexpr std::uint64_t low() const { return value_.low(); }

    template <class T, std::enable_if_t<std::is_integral_v<T>, int> = 0>
    explicit constexpr operator T() const {
        return static_cast<T>(value_.low());
    }

    explicit constexpr operator bool() const { return value_ != UInt128{}; }

    constexpr bool is_negative() const { return (value_.high() >> 63) != 0; }

    explicit constexpr operator long double() const {
        const bool negative = is_negative();
        const long double magnitude =
            static_cast<long double>(negative ? -value_ : value_);
        return negative ? -magnitude : magnitude;
    }

    friend constexpr bool operator==(Int128 lhs, Int128 rhs) {
        return lhs.value_ == rhs.value_;
    }

    friend constexpr bool operator!=(Int128 lhs, Int128 rhs) {
        return !(lhs == rhs);
    }

    friend constexpr bool operator<(Int128 lhs, Int128 rhs) {
        const bool lhs_negative = lhs.is_negative();
        const bool rhs_negative = rhs.is_negative();
        if (lhs_negative != rhs_negative) {
            return lhs_negative;
        }
        return lhs.value_ < rhs.value_;
    }

    friend constexpr bool operator>(Int128 lhs, Int128 rhs) {
        return rhs < lhs;
    }

    friend constexpr bool operator<=(Int128 lhs, Int128 rhs) {
        return !(rhs < lhs);
    }

    friend constexpr bool operator>=(Int128 lhs, Int128 rhs) {
        return !(lhs < rhs);
    }

    constexpr Int128 operator+() const { return *this; }

    constexpr Int128 operator-() const {
        Int128 value;
        value.value_ = -value_;
        return value;
    }

    constexpr Int128 &operator+=(Int128 rhs) {
        value_ += rhs.value_;
        return *this;
    }

    constexpr Int128 &operator-=(Int128 rhs) {
        value_ -= rhs.value_;
        return *this;
    }

    constexpr Int128 &operator*=(Int128 rhs) {
        value_ *= rhs.value_;
        return *this;
    }

    constexpr Int128 &operator/=(Int128 rhs) {
        *this = *this / rhs;
        return *this;
    }

    constexpr Int128 &operator%=(Int128 rhs) {
        *this = *this % rhs;
        return *this;
    }

    friend constexpr Int128 operator+(Int128 lhs, Int128 rhs) {
        lhs += rhs;
        return lhs;
    }

    friend constexpr Int128 operator-(Int128 lhs, Int128 rhs) {
        lhs -= rhs;
        return lhs;
    }

    friend constexpr Int128 operator*(Int128 lhs, Int128 rhs) {
        lhs *= rhs;
        return lhs;
    }

    static constexpr void div_mod(Int128 lhs, Int128 rhs, Int128 &quotient,
                                  Int128 &remainder) {
        assert(rhs != Int128{});
        assert(!(lhs == min_value() && rhs == Int128(-1)));

        const bool lhs_negative = lhs.is_negative();
        const bool rhs_negative = rhs.is_negative();
        const bool quotient_negative = lhs_negative != rhs_negative;

        UInt128 quotient_abs;
        UInt128 remainder_abs;
        UInt128::div_mod_unchecked(abs_unsigned(lhs), abs_unsigned(rhs),
                                   quotient_abs, remainder_abs);

        const Int128 q =
            from_unsigned_unchecked(quotient_abs, quotient_negative);
        const Int128 r = from_unsigned_unchecked(remainder_abs, lhs_negative);

        quotient = q;
        remainder = r;
    }

    friend constexpr Int128 operator/(Int128 lhs, Int128 rhs) {
        Int128 quotient;
        Int128 remainder;
        div_mod(lhs, rhs, quotient, remainder);
        return quotient;
    }

    friend constexpr Int128 operator%(Int128 lhs, Int128 rhs) {
        Int128 quotient;
        Int128 remainder;
        div_mod(lhs, rhs, quotient, remainder);
        return remainder;
    }

  private:
    static constexpr Int128 from_twos_complement(UInt128 value) {
        Int128 result;
        result.value_ = value;
        return result;
    }

    static constexpr Int128 min_value() {
        return from_words(std::uint64_t{1} << 63, 0);
    }

    static constexpr UInt128 abs_unsigned(Int128 value) {
        if (!value.is_negative()) {
            return value.value_;
        }
        return -value.value_;
    }

    static constexpr Int128 from_unsigned(UInt128 value, bool negative) {
        if (!negative) {
            assert((value.high() >> 63) == 0);
        } else {
            assert(value <= UInt128::from_words(std::uint64_t{1} << 63, 0));
        }
        return from_unsigned_unchecked(value, negative);
    }

    // 引数 negative で指定した通りに符号を付けた value が
    // Int128 の範囲に収まることを仮定する。
    static constexpr Int128 from_unsigned_unchecked(UInt128 value,
                                                    bool negative) {
        return from_twos_complement(negative ? -value : value);
    }

    friend std::istream &operator>>(std::istream &input, Int128 &value);
    friend std::ostream &operator<<(std::ostream &output, Int128 value);

    UInt128 value_;
};

namespace int128_internal {
// first < text.size() が呼び出し元で保証されていることを仮定する。
inline UInt128 read_uint128_decimal(const std::string &text,
                                    std::size_t first) {
    UInt128 value = 0;

    std::size_t i = first;
    const std::size_t first_len = (text.size() - first) % 9;
    if (first_len != 0) {
        std::uint32_t chunk = 0;
        for (std::size_t j = 0; j < first_len; ++j) {
            assert('0' <= text[i] && text[i] <= '9');
            chunk = chunk * 10 + static_cast<std::uint32_t>(text[i] - '0');
            ++i;
        }
        value = UInt128(chunk);
    }

    while (i < text.size()) {
        std::uint32_t chunk = 0;
        for (int j = 0; j < 9; ++j) {
            assert('0' <= text[i] && text[i] <= '9');
            chunk = chunk * 10 + static_cast<std::uint32_t>(text[i] - '0');
            ++i;
        }
        value *= UInt128(1000000000);
        value += UInt128(chunk);
    }

    return value;
}

inline void write_padded_9(std::ostream &output, std::uint32_t value) {
    char buffer[9];
    for (int i = 8; i >= 0; --i) {
        buffer[i] = static_cast<char>('0' + value % 10);
        value /= 10;
    }
    output.write(buffer, 9);
}

inline void write_uint128_decimal(std::ostream &output, UInt128 value) {
    if (value == UInt128(0)) {
        output << '0';
        return;
    }

    constexpr std::uint32_t base = 1000000000;
    std::uint32_t parts[5] = {};
    int size = 0;
    while (value != UInt128(0)) {
        UInt128 quotient;
        UInt128 remainder;
        UInt128::div_mod(value, UInt128(base), quotient, remainder);
        parts[size] = static_cast<std::uint32_t>(remainder);
        value = quotient;
        ++size;
    }

    output << parts[size - 1];
    for (int i = size - 2; i >= 0; --i) {
        write_padded_9(output, parts[i]);
    }
}
} // namespace int128_internal

inline std::istream &operator>>(std::istream &input, UInt128 &value) {
    std::string text;
    input >> text;
    if (!input) {
        return input;
    }

    std::size_t first = 0;
    if (text[first] == '+') {
        ++first;
    }
    assert(first < text.size());
    assert(text[first] != '-');

    value = int128_internal::read_uint128_decimal(text, first);
    return input;
}

inline std::ostream &operator<<(std::ostream &output, UInt128 value) {
    int128_internal::write_uint128_decimal(output, value);
    return output;
}

inline std::istream &operator>>(std::istream &input, Int128 &value) {
    std::string text;
    input >> text;
    if (!input) {
        return input;
    }

    std::size_t first = 0;
    bool negative = false;
    if (text[first] == '+' || text[first] == '-') {
        negative = text[first] == '-';
        ++first;
    }
    assert(first < text.size());

    const UInt128 magnitude =
        int128_internal::read_uint128_decimal(text, first);
    value = Int128::from_unsigned(magnitude, negative);
    return input;
}

inline std::ostream &operator<<(std::ostream &output, Int128 value) {
    if (value.is_negative()) {
        output << '-';
        int128_internal::write_uint128_decimal(output,
                                               Int128::abs_unsigned(value));
    } else {
        int128_internal::write_uint128_decimal(
            output, UInt128::from_words(value.high(), value.low()));
    }
    return output;
}
} // namespace NicheLibrary

namespace std {
template <> class numeric_limits<NicheLibrary::UInt128> {
  public:
    static constexpr bool is_specialized = true;

    static constexpr NicheLibrary::UInt128 min() noexcept { return 0; }
    static constexpr NicheLibrary::UInt128 max() noexcept {
        return NicheLibrary::UInt128::from_words(~std::uint64_t{},
                                                 ~std::uint64_t{});
    }
    static constexpr NicheLibrary::UInt128 lowest() noexcept { return 0; }

    static constexpr int digits = 128;
    static constexpr int digits10 = 38;
    static constexpr int max_digits10 = 0;
    static constexpr bool is_signed = false;
    static constexpr bool is_integer = true;
    static constexpr bool is_exact = true;
    static constexpr int radix = 2;

    static constexpr NicheLibrary::UInt128 epsilon() noexcept { return 0; }
    static constexpr NicheLibrary::UInt128 round_error() noexcept { return 0; }

    static constexpr int min_exponent = 0;
    static constexpr int min_exponent10 = 0;
    static constexpr int max_exponent = 0;
    static constexpr int max_exponent10 = 0;

    static constexpr bool has_infinity = false;
    static constexpr bool has_quiet_NaN = false;
    static constexpr bool has_signaling_NaN = false;
    static constexpr float_denorm_style has_denorm = denorm_absent;
    static constexpr bool has_denorm_loss = false;

    static constexpr NicheLibrary::UInt128 infinity() noexcept { return 0; }
    static constexpr NicheLibrary::UInt128 quiet_NaN() noexcept { return 0; }
    static constexpr NicheLibrary::UInt128 signaling_NaN() noexcept {
        return 0;
    }
    static constexpr NicheLibrary::UInt128 denorm_min() noexcept { return 0; }

    static constexpr bool is_iec559 = false;
    static constexpr bool is_bounded = true;
    static constexpr bool is_modulo = true;
    static constexpr bool traps = false;
    static constexpr bool tinyness_before = false;
    static constexpr float_round_style round_style = round_toward_zero;
};

template <> class numeric_limits<NicheLibrary::Int128> {
  public:
    static constexpr bool is_specialized = true;

    static constexpr NicheLibrary::Int128 min() noexcept {
        return NicheLibrary::Int128::from_words(std::uint64_t{1} << 63, 0);
    }
    static constexpr NicheLibrary::Int128 max() noexcept {
        return NicheLibrary::Int128::from_words((std::uint64_t{1} << 63) - 1,
                                                ~std::uint64_t{});
    }
    static constexpr NicheLibrary::Int128 lowest() noexcept { return min(); }

    static constexpr int digits = 127;
    static constexpr int digits10 = 38;
    static constexpr int max_digits10 = 0;
    static constexpr bool is_signed = true;
    static constexpr bool is_integer = true;
    static constexpr bool is_exact = true;
    static constexpr int radix = 2;

    static constexpr NicheLibrary::Int128 epsilon() noexcept { return 0; }
    static constexpr NicheLibrary::Int128 round_error() noexcept { return 0; }

    static constexpr int min_exponent = 0;
    static constexpr int min_exponent10 = 0;
    static constexpr int max_exponent = 0;
    static constexpr int max_exponent10 = 0;

    static constexpr bool has_infinity = false;
    static constexpr bool has_quiet_NaN = false;
    static constexpr bool has_signaling_NaN = false;
    static constexpr float_denorm_style has_denorm = denorm_absent;
    static constexpr bool has_denorm_loss = false;

    static constexpr NicheLibrary::Int128 infinity() noexcept { return 0; }
    static constexpr NicheLibrary::Int128 quiet_NaN() noexcept { return 0; }
    static constexpr NicheLibrary::Int128 signaling_NaN() noexcept { return 0; }
    static constexpr NicheLibrary::Int128 denorm_min() noexcept { return 0; }

    static constexpr bool is_iec559 = false;
    static constexpr bool is_bounded = true;
    static constexpr bool is_modulo = false;
    static constexpr bool traps = false;
    static constexpr bool tinyness_before = false;
    static constexpr float_round_style round_style = round_toward_zero;
};
} // namespace std


#line 17 "math/number-theory/generalized-floor-sum-degree-le-2.hpp"

template <class T> struct GeneralizedFloorSumDegreeLe2Result {
    T ans_01;
    T ans_11;
    T ans_02;
};

namespace generalized_floor_sum_degree_le_2_internal {
template <class T>
constexpr bool is_integer_v =
    std::numeric_limits<std::remove_cv_t<T>>::is_integer;

template <class T>
constexpr bool is_signed_v =
    std::numeric_limits<std::remove_cv_t<T>>::is_signed;

template <class T>
constexpr bool use_int128_by_default_v =
    std::is_integral_v<std::remove_cv_t<T>> &&
    !std::is_same_v<std::remove_cv_t<T>, bool> &&
    std::numeric_limits<std::remove_cv_t<T>>::digits <= 64;

template <class T>
using default_internal_t =
    std::conditional_t<use_int128_by_default_v<T>, NicheLibrary::Int128, T>;

template <class T> std::pair<T, T> floor_div_mod(T x, T y) {
    if constexpr (is_signed_v<T>) {
        if (x >= 0 && x < y) {
            return {0, x};
        }
        T q = x / y;
        T r = x - q * y;
        if (r < 0) {
            q -= T(1);
            r += y;
        }
        return {q, r};
    } else {
        if (x < y) {
            return {0, x};
        }
        const T q = x / y;
        return {q, x - q * y};
    }
}

// Σ_{i=0}^{n-1} i = n(n-1)/2
template <class State, class Value> Value sum_0_to_n_minus_1(State n) {
    const State half = n / 2;
    if (n == half * 2) {
        return static_cast<Value>(half) * static_cast<Value>(n - 1);
    }
    return static_cast<Value>(n) * static_cast<Value>(half);
}

// Σ_{i=0}^{n-1} i^2 = (n-1)n(2n-1)/6
template <class State, class Value> Value sum_0_to_n_minus_1_sq(State n) {
    Value a = static_cast<Value>(n - 1);
    Value b = static_cast<Value>(n);
    Value c = static_cast<Value>(2) * static_cast<Value>(n) - Value(1);
    const State n_mod_6 = n % 6;
    if (n_mod_6 == 0) {
        b /= Value(6);
    } else if (n_mod_6 == 1) {
        a /= Value(6);
    } else if (n_mod_6 == 2) {
        b /= Value(2);
        c /= Value(3);
    } else if (n_mod_6 == 3) {
        a /= Value(2);
        b /= Value(3);
    } else if (n_mod_6 == 4) {
        a /= Value(3);
        b /= Value(2);
    } else {
        a /= Value(2);
        c /= Value(3);
    }
    return a * b * c;
}

// Σ_{i=l}^{n-1} i
template <class State, class Value> Value sum_l_to_n_minus_1(State l, State n) {
    const State count = n - l;
    const Value sum = static_cast<Value>(l) + static_cast<Value>(n) - Value(1);
    const State half_count = count / 2;
    if (count == half_count * 2) {
        return static_cast<Value>(half_count) * sum;
    }
    return static_cast<Value>(count) * (sum / Value(2));
}

template <class Value> struct Result {
    Value ans_01;
    Value ans_11;
    Value ans_02;
};

template <class State, class Value>
Result<Value> solve(State n, State m, State a, State b) {
    const auto [qa_state, a_mod] = floor_div_mod(a, m);
    const auto [qb_state, b_mod] = floor_div_mod(b, m);
    a = a_mod;
    b = b_mod;

    const Value qa = static_cast<Value>(qa_state);
    const Value qb = static_cast<Value>(qb_state);

    Result<Value> base = {0, 0, 0};

    if (a != 0) {
        const Value y_max_value =
            (static_cast<Value>(a) * static_cast<Value>(n) +
             static_cast<Value>(b)) /
            static_cast<Value>(m);

        const State y_max = static_cast<State>(y_max_value);

        if (y_max != 0) {
            const Value x_max =
                y_max_value * static_cast<Value>(m) - static_cast<Value>(b);

            const Value a_value = static_cast<Value>(a);
            const Value x_max_div_a = x_max / a_value;
            const Value x_max_mod_a = x_max - x_max_div_a * a_value;
            const bool has_remainder = x_max_mod_a != Value(0);
            const State t = static_cast<State>(x_max_div_a) +
                            (has_remainder ? State(1) : State(0));
            const State b2 = has_remainder
                                 ? static_cast<State>(a_value - x_max_mod_a)
                                 : State(0);

            const auto rec = solve<State, Value>(y_max, a, m, b2);

            const Value head_01 = rec.ans_01;
            const Value head_11 =
                ((static_cast<Value>(2) * static_cast<Value>(t) - Value(1)) *
                     rec.ans_01 -
                 rec.ans_02) /
                Value(2);
            const Value head_02 =
                (static_cast<Value>(2) * static_cast<Value>(y_max) - Value(1)) *
                    rec.ans_01 -
                static_cast<Value>(2) * rec.ans_11;

            const State tail_len_state = n - t;
            const Value tail_len = static_cast<Value>(tail_len_state);
            const Value y = static_cast<Value>(y_max);

            const Value tail_01 = tail_len * y;
            const Value tail_11 = y * sum_l_to_n_minus_1<State, Value>(t, n);
            const Value tail_02 = tail_len * y * y;

            base = {head_01 + tail_01, head_11 + tail_11, head_02 + tail_02};
        }
    }

    const Value n_value = static_cast<Value>(n);
    if (qa == Value(0)) {
        if (qb == Value(0)) {
            return base;
        }
        const Value si = sum_0_to_n_minus_1<State, Value>(n);
        return {qb * n_value + base.ans_01, qb * si + base.ans_11,
                qb * qb * n_value + Value(2) * qb * base.ans_01 + base.ans_02};
    }

    const Value si = sum_0_to_n_minus_1<State, Value>(n);
    const Value si2 = sum_0_to_n_minus_1_sq<State, Value>(n);
    if (qb == Value(0)) {
        return {qa * si + base.ans_01, qa * si2 + base.ans_11,
                qa * qa * si2 + Value(2) * qa * base.ans_11 + base.ans_02};
    }

    Result<Value> ans;
    ans.ans_01 = qa * si + qb * n_value + base.ans_01;
    ans.ans_11 = qa * si2 + qb * si + base.ans_11;
    ans.ans_02 = qa * qa * si2 + Value(2) * qa * qb * si +
                 Value(2) * qa * base.ans_11 + qb * qb * n_value +
                 Value(2) * qb * base.ans_01 + base.ans_02;
    return ans;
}
} // namespace generalized_floor_sum_degree_le_2_internal

template <class T, class Internal = void>
GeneralizedFloorSumDegreeLe2Result<T>
generalized_floor_sum_degree_le_2(T n, T m, T a, T b) {
    namespace gfs_internal = generalized_floor_sum_degree_le_2_internal;
    using Value =
        std::conditional_t<std::is_void_v<Internal>,
                           gfs_internal::default_internal_t<T>, Internal>;

    static_assert(gfs_internal::is_integer_v<T>, "T must be integer.");
    static_assert(gfs_internal::is_integer_v<Value>,
                  "Internal must be integer.");
    static_assert(!std::is_same_v<std::remove_cv_t<T>, bool>,
                  "T must not be bool.");
    static_assert(!std::is_same_v<std::remove_cv_t<Value>, bool>,
                  "Internal must not be bool.");
    static_assert(!gfs_internal::is_signed_v<T> ||
                      gfs_internal::is_signed_v<Value>,
                  "Internal must be signed when T is signed.");
    static_assert(std::numeric_limits<std::remove_cv_t<T>>::digits <=
                      std::numeric_limits<std::remove_cv_t<Value>>::digits,
                  "Internal must be able to represent values of T.");

    if constexpr (gfs_internal::is_signed_v<T>) {
        assert(n >= 0);
    }
    assert(m > 0);
    if (n == 0) {
        return {T(0), T(0), T(0)};
    }
    if (n == 1) {
        const T q = gfs_internal::floor_div_mod(b, m).first;
        const Value q_value = static_cast<Value>(q);
        return {q, T(0), static_cast<T>(q_value * q_value)};
    }

    const auto res = gfs_internal::solve<T, Value>(n, m, a, b);
    return {static_cast<T>(res.ans_01), static_cast<T>(res.ans_11),
            static_cast<T>(res.ans_02)};
}


#line 7 "verify/standalone-generalized-floor-sum-degree-le-2.test.cpp"

long long floor_div_brute(long long x, long long y) {
    assert(y > 0);

    long long q = x / y;
    long long r = x % y;
    if (r < 0) {
        --q;
    }
    return q;
}

int main() {
    for (long long n = 0; n <= 20; ++n) {
        for (long long m = 1; m <= 20; ++m) {
            for (long long a = -20; a <= 20; ++a) {
                for (long long b = -20; b <= 20; ++b) {
                    const auto res =
                        generalized_floor_sum_degree_le_2<long long>(n, m, a,
                                                                     b);

                    long long s_01 = 0, s_11 = 0, s_02 = 0;
                    for (long long i = 0; i < n; ++i) {
                        const long long value = floor_div_brute(a * i + b, m);
                        s_01 += value;
                        s_11 += i * value;
                        s_02 += value * value;
                    }

                    assert(res.ans_01 == s_01);
                    assert(res.ans_11 == s_11);
                    assert(res.ans_02 == s_02);
                }
            }
        }
    }

    {
        using T = std::uint64_t;

        const T n = 900931385;
        const T m = 333006410;
        const T x = 263208878;
        const T y = 243209245;

        const auto r0 = generalized_floor_sum_degree_le_2<T>(n, m, x, T(0));
        const auto r1 = generalized_floor_sum_degree_le_2<T>(n, m, x, y);

        T ans = 0;
        ans -= r0.ans_01 * n - r0.ans_11;
        ans -= r1.ans_01 * (n - 1) - 2 * r1.ans_11;

        assert(ans == 202919340569980512ULL);
    }

    return 0;
}

Back to top page