Number_Theory/SumofFloor.hpp

• View this file on GitHub
• Last update: 2023-02-11 22:28:05+08:00
• Include: #include "Number_Theory/SumofFloor.hpp"

Depends on

• Template/Template.hpp

Verified with

• Verify/SumofFloorofLinear.test.cpp

Code

#pragma once

#include "../Template/Template.hpp"
// sum_{0 <= i < N} (ai + b) // m
template <typename T>
T sum_of_floor(T n, T m, T a, T b) {
    T ret = 0;
    if (a >= m) ret += (n - 1) * n * (a / m) / 2, a %= m;
    if (b >= m) ret += n * (b / m), b %= m;
    T y = (a * n + b) / m;
    if (y == 0) return ret;
    T x = y * m - b;
    ret += (n - (x + a - 1) / a) * y;
    ret += sum_of_floor(y, a, m, (a - x % a) % a);
    return ret;
}

// verify www.codechef.com/viewsolution/36222026
// count x : ax + b mod m < yr, 0 <= x < xr
template <typename T>
T mod_affine_range_counting(T a, T b, T m, T xr, T yr) {
    assert(0 <= yr && yr <= m);
    return sum_of_floor(xr, m, a, b + m) - sum_of_floor(xr, m, a, b + m - yr);
}

#line 2 "Number_Theory/SumofFloor.hpp"

#line 2 "Template/Template.hpp"

#include <bits/stdc++.h>

using i64 = std::int64_t;
#line 4 "Number_Theory/SumofFloor.hpp"
// sum_{0 <= i < N} (ai + b) // m
template <typename T>
T sum_of_floor(T n, T m, T a, T b) {
    T ret = 0;
    if (a >= m) ret += (n - 1) * n * (a / m) / 2, a %= m;
    if (b >= m) ret += n * (b / m), b %= m;
    T y = (a * n + b) / m;
    if (y == 0) return ret;
    T x = y * m - b;
    ret += (n - (x + a - 1) / a) * y;
    ret += sum_of_floor(y, a, m, (a - x % a) % a);
    return ret;
}

// verify www.codechef.com/viewsolution/36222026
// count x : ax + b mod m < yr, 0 <= x < xr
template <typename T>
T mod_affine_range_counting(T a, T b, T m, T xr, T yr) {
    assert(0 <= yr && yr <= m);
    return sum_of_floor(xr, m, a, b + m) - sum_of_floor(xr, m, a, b + m - yr);
}

Back to top page