math/number/floor-sum.hpp

• View this file on GitHub
• Last update: 2024-07-17 12:55:05+09:00
• Include: #include "math/number/floor-sum.hpp"

Required by

• scrap/static-top-tree.hpp

Verified with

• test/yosupo/number_theory/sum_of_floor_of_linear.test.cpp

Code

#pragma once

template <typename T>
T floor_sum(T n, T m, T a, T b) {
    T ans = 0;
    if (a < 0) {
        const T a2 = a % m + m;
        ans -= n * (n - 1) / 2 * ((a2 - a) / m);
        a = a2;
    }
    if (b < 0) {
        const T b2 = b % m + m;
        ans -= n * ((b2 - b) / m);
        b = b2;
    }
    while (true) {
        if (a >= m) {
            ans += n * (n - 1) / 2 * (a / m);
            a %= m;
        }
        if (b >= m) {
            ans += 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 ans;
}

#line 2 "math/number/floor-sum.hpp"

template <typename T>
T floor_sum(T n, T m, T a, T b) {
    T ans = 0;
    if (a < 0) {
        const T a2 = a % m + m;
        ans -= n * (n - 1) / 2 * ((a2 - a) / m);
        a = a2;
    }
    if (b < 0) {
        const T b2 = b % m + m;
        ans -= n * ((b2 - b) / m);
        b = b2;
    }
    while (true) {
        if (a >= m) {
            ans += n * (n - 1) / 2 * (a / m);
            a %= m;
        }
        if (b >= m) {
            ans += 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 ans;
}

Back to top page