sotanishy's code snippets for competitive programming
#include "math/number-theory/floor_sum.hpp"
一次関数の床関数の和 $\sum_{i=0}^{N-1} \left\lfloor \frac{Ai + B}{M} \right\rfloor$ を再帰的に計算する.
計算量はユークリッドの互除法と同様に対数時間となるが,どの変数に依存するのかはよくわかっていない (僕が理解していない).
long long floor_sum(long long n, long long m, long long a, long long b)
#pragma once
long long floor_sum(long long n, long long m, long long a, long long b) {
long long sum = 0;
if (a >= m) {
sum += (a / m) * n * (n - 1) / 2;
a %= m;
}
if (b >= m) {
sum += (b / m) * n;
b %= m;
}
long long y = (a * n + b) / m;
if (y == 0) return sum;
long long x = (m * y - b + a - 1) / a;
sum += (n - x) * y + floor_sum(y, a, m, a * x - m * y + b);
return sum;
}
#line 2 "math/number-theory/floor_sum.hpp"
long long floor_sum(long long n, long long m, long long a, long long b) {
long long sum = 0;
if (a >= m) {
sum += (a / m) * n * (n - 1) / 2;
a %= m;
}
if (b >= m) {
sum += (b / m) * n;
b %= m;
}
long long y = (a * n + b) / m;
if (y == 0) return sum;
long long x = (m * y - b + a - 1) / a;
sum += (n - x) * y + floor_sum(y, a, m, a * x - m * y + b);
return sum;
}