Floor Sum
(math/floor-sum.hpp)
- View this file on GitHub
- Last update: 2025-10-17 21:43:09+09:00
- Include:
#include "math/floor-sum.hpp"
使い方
以下の問題を解く.
整数 $n,m,a,b\ (m\neq 0)$ に対し以下の値を求めよ: \(f(n,m,a,b)=\sum_{i=0}^{n-1}\left\lfloor\frac{ai+b}{m}\right\rfloor\)
$an+b$ が unsigned long long におさまる必要がある.
アルゴリズム
再帰的に考える.
$n=0$ のとき $f(n,m,a,b)=0$ である.
簡単な変形によって $m\geq 1$ および $0\leq a\lt m, 0\leq b\lt m$ として構わない.
$a\geq 0$ のとき $an+b=mq+r\,(0\leq r\lt m)$ とする. $0\leq i\lt n$ に対し $0\leq\left\lfloor\frac{ai+b}{m}\right\rfloor\leq q$ であり,$\left\lfloor\frac{ai+b}{m}\right\rfloor\geq k$ となる $i$ の範囲は $\frac{mk-b}{a}\leq i$ であるから以下が成り立つ.
\[\begin{align*} f(n,m,a,b) &=\sum_{k=1}^{q}\left(n-\left\lceil\frac{mk-b}{a}\right\rceil\right)\\ &=\sum_{k=1}^{q}\left(n+\left\lfloor\frac{b-mk}{a}\right\rfloor\right)\\ &=\sum_{k=0}^{q-1}\left(n+\left\lfloor\frac{b-m(q-k)}{a}\right\rfloor\right)\\ &=\sum_{k=0}^{q-1}\left\lfloor\frac{mk+an+b-mq}{a}\right\rfloor\\ &=\sum_{k=0}^{q-1}\left\lfloor\frac{mk+r}{a}\right\rfloor\\ &=f(q,a,m,r) \end{align*}\]$m,a$ について Euclid の互除法と同じ議論ができ,$O(\log m)$ 回の反復で停止することが示せる.
Verified with
Code
#pragma once
// sum{i=0}^{n-1}floor((a*i+b)/m)
template <class T>
T FloorSumUnsigned(unsigned long long n, unsigned long long m, unsigned long long a, unsigned long long b) {
assert(m != 0);
if (n == 0) return 0;
T res = 0;
while (true) {
if (a >= m) {
unsigned long long q = a / m;
res += T(q) * (n / 2) * ((n - 1) | 1);
a -= m * q;
}
if (b >= m) {
unsigned long long q = b / m;
res += T(q) * n;
b -= m * q;
}
unsigned long long y = a * n + b;
if (y < m) break;
n = y / m, b = y - m * n;
swap(a, m);
}
return res;
}
// sum{i=0}^{n-1}floor((a*i+b)/m)
template <class T>
T FloorSum(long long n, long long m, long long a, long long b) {
assert(m != 0);
if (n <= 0) return 0;
if (m < 0) a = -a, b = -b, m = -m;
T res = 0;
if (a < 0) {
long long q = (a + 1) / m - 1;
res += T(q) * (n / 2) * ((n - 1) | 1);
a -= m * q;
}
if (b < 0) {
long long q = (b + 1) / m - 1;
res += T(q) * n;
b -= m * q;
}
return res + FloorSumUnsigned<T>(n, m, a, b);
}
/**
* @brief Floor Sum
* @docs docs/math/floor-sum.md
*/#line 2 "math/floor-sum.hpp"
// sum{i=0}^{n-1}floor((a*i+b)/m)
template <class T>
T FloorSumUnsigned(unsigned long long n, unsigned long long m, unsigned long long a, unsigned long long b) {
assert(m != 0);
if (n == 0) return 0;
T res = 0;
while (true) {
if (a >= m) {
unsigned long long q = a / m;
res += T(q) * (n / 2) * ((n - 1) | 1);
a -= m * q;
}
if (b >= m) {
unsigned long long q = b / m;
res += T(q) * n;
b -= m * q;
}
unsigned long long y = a * n + b;
if (y < m) break;
n = y / m, b = y - m * n;
swap(a, m);
}
return res;
}
// sum{i=0}^{n-1}floor((a*i+b)/m)
template <class T>
T FloorSum(long long n, long long m, long long a, long long b) {
assert(m != 0);
if (n <= 0) return 0;
if (m < 0) a = -a, b = -b, m = -m;
T res = 0;
if (a < 0) {
long long q = (a + 1) / m - 1;
res += T(q) * (n / 2) * ((n - 1) | 1);
a -= m * q;
}
if (b < 0) {
long long q = (b + 1) / m - 1;
res += T(q) * n;
b -= m * q;
}
return res + FloorSumUnsigned<T>(n, m, a, b);
}
/**
* @brief Floor Sum
* @docs docs/math/floor-sum.md
*/