floor sum アルゴリズムの導出方法とその一般化について紹介します。
floor sum アルゴリズムとは
$\displaystyle f(N,M,A,B)=\sum_{i=0}^{N-1}\left\lfloor\frac{Ai+B}{M}\right\rfloor$ と定義します。
整数 $N,M,A,B$ が与えられた時に $f(N,M,A,B)$ を計算量 $O(\log M)$ で求めるアルゴリズムが floor sum アルゴリズムです。
記号・式の紹介
導出の際に必要となる記号や式についてまとめておきます。ただし、 $x$ は実数、 $X,Y$ は整数、 $Y\neq 0$ とします。
- $\lfloor x\rfloor =(x$ 以下の最大の整数$)$ と定義する
- $\lceil x\rceil = (x$ 以上の最小の整数$)$ と定義する
- $X=X_1M+X_2$ $(0\le X_2 <M)$ としたとき、
- $\displaystyle X_1=\left\lfloor \frac{X}M\right\rfloor$
- $\displaystyle X_2=X\ \text{mod}\ M$
- $\displaystyle\left\lceil \frac XY\right\rceil= \left\lfloor \frac {X+Y-1}Y\right\rfloor$
- $\left\lfloor X+x\right\rfloor = X + \left\lfloor x\right\rfloor$
- $x < X\iff \left\lfloor x\right\rfloor < X$
- $x\le X\iff \left\lceil x\right\rceil \le X$
- $x> X\iff \left\lceil x\right\rceil > X$
- $x\geq X \iff \left\lfloor x\right\rfloor\geq X$
導出方法
$A=A_1M+A_2$ とすると、
$
\begin{align*}
f(N,M,A,B) &= \sum_{i=0}^{N-1}\left\lfloor\frac{Ai+B}{M}\right\rfloor \\
&=\sum_{i=0}^{N-1}\left\lfloor\frac{(A_1M+A_2)i+B}{M}\right\rfloor \\
&=\sum_{i=0}^{N-1}\left\lfloor A_1i + \frac{A_2i+B}{M}\right\rfloor \\
&=\sum_{i=0}^{N-1} \left(A_1i +\left\lfloor \frac{A_2i+B}{M}\right\rfloor\right) \\
&=A_1\frac{N(N-1)}2+ \sum_{i=0}^{N-1} \left\lfloor \frac{A_2i+B}{M}\right\rfloor \\
&=A_1\frac{N(N-1)}2+f(N,M,A_2,B)
\end{align*}
$
が成立します。また、同様に $B=B_1M+B_2$ とすると、
$
\begin{align*}
f(N,M,A,B) &= \sum_{i=0}^{N-1}\left\lfloor\frac{Ai+B}{M}\right\rfloor \\
&=\sum_{i=0}^{N-1}\left\lfloor\frac{Ai+B_1M+B_2}{M}\right\rfloor \\
&=\sum_{i=0}^{N-1}\left\lfloor B_1+\frac{Ai+B_2}{M}\right\rfloor \\
&=\sum_{i=0}^{N-1}\left(B_1+\left\lfloor \frac{Ai+B_2}{M}\right\rfloor \right)\\
&=B_1N+\sum_{i=0}^{N-1} \left\lfloor \frac{Ai+B_2}{M}\right\rfloor \\
&=B_1N+f(N,M,A,B_2)
\end{align*}
$
が成立します。
この $2$ つの式を用いることで、引数 $A,B$ の値をどちらも $0$ 以上 $M$ 未満の場合に帰着させることができます。
ここで、 $\displaystyle K=\left\lfloor\frac{A(N-1)+B}{M}\right\rfloor$ と定義すると、 $0\le i\le N - 1$ に対して $\displaystyle \left\lfloor\frac{Ai+B}{M}\right\rfloor$ の取る値は $0,1,\ldots,K$ のいずれかになります。よって、 $\displaystyle \left\lfloor\frac{Ai+B}{M}\right\rfloor=k$ を満たす $0\le i\le N - 1$ の個数を $c_k$ とすると、 $\displaystyle f(N,M,A,B)=\sum_{k=0}^K kc_k$ が成立することが分かります。(主客転倒)
ここで、条件 $\displaystyle \left\lfloor\frac{Ai+B}{M}\right\rfloor=k$ を同値変形していきます。
$
\begin{align*}
& \left\lfloor\frac{Ai+B}{M}\right\rfloor=k\\
\iff & k\le \frac{Ai+B}{M} < k + 1\\
\iff &\frac{kM-B}A\le i < \frac{(k+1)M-B}A\\
\iff &\left\lceil\frac{kM-B}A \right\rceil \le i < \left\lceil\frac{(k+1)M-B}A \right\rceil
\end{align*}
$
よって、 $\displaystyle d_k=\left\lceil\frac{kM-B}A \right\rceil$ とすると、 $\displaystyle \left\lfloor\frac{Ai+B}{M}\right\rfloor=k \iff d_k\le i < d_{k+1}$ となります。ただし、 $k=K$ の場合は $i=N-1$ が上限なので $d_{K+1}=N$ とします。
以上より $c_k=d_{k+1}-d_k$ が分かるので、
$
\begin{align*}
f(N,M,A,B)&= \sum_{k=0}^K kc_k\\
&= \sum_{k=0}^K k(-d_k+d_{k+1})\\
&= -d_1+(1-2)d_2+(2-3)d_3+\cdots+((K-1)-K)d_K+Kd_{K+1}\\
&=KN-\sum_{k=1}^K d_k\\
&=KN-\sum_{k=1}^K \left\lceil\frac{kM-B}A \right\rceil\\
&=KN-\sum_{k=1}^K \left\lfloor\frac{kM-B+A-1}A \right\rfloor\\
&=KN-\sum_{k=0}^{K-1} \left\lfloor\frac{(k+1)M-B+A-1}A \right\rfloor\\
&=KN-\sum_{k=0}^{K-1} \left\lfloor\frac{kM+M-B+A-1}A \right\rfloor\\
&=KN-f(K,A,M,M-B+A-1)
\end{align*}
$
が従います。つまり、 $f(N,M,A,B)$ を求めるためには $f(K,A,M,M-B+A-1)$ を求めれば良いということです。
この $f$ の第二引数と第三引数に注目します。
$f(N,M,A,B)$ を求めるためには $f(N,A,M,M-B+A-1)$ を求めれば良く、さらにそれを求めるためには $f(N,A,M\ \text{mod} \ A,M-B+A-1)$ を求めれば良いです。この際の第二引数と第三引数の組は $(M,A), (A,M) , (A, M\ \text{mod}\ A),\ldots$ と移っていきますが、これはまさにユークリッドの互除法で $\gcd(M,A)$ を求める際の引数のペアの遷移と同じです。
$\gcd(M,A)$を求める際のステップ数は $O(\log M)$ である (ラメの定理) ので、この floor sum の再帰の回数も $O(\log M)$ 回であるということが分かります。
floor sum アルゴリズムの終了条件は $A=0$ で、その際の値は $\displaystyle f(N,M,0,B)=\sum_{i=0}^{N-1}\left\lfloor\frac{0\times i+B}{M}\right\rfloor=N\left\lfloor\frac{B}{M} \right\rfloor$ です。
Python で書くと以下のようになります。
def floor_sum(n, m, a, b):
a1, a2 = a // m, a % m
s = n * (n - 1) // 2 * a1
b1, b2 = b // m, b % m
if a2 == 0:
return s + b1 * n
k = (a2 * (n - 1) + b2) // m
return s + n * (k + b1) - floor_sum(k, a2, m, m + a2 - b2 - 1)
例題
Min of Mod of Linear(yosupo judge)
整数 $N,M,A,B$ が与えられるので、 $\displaystyle \min_{0\le x<N}\lbrace (Ax+B)\ \text{mod}\ M\rbrace $ を求めてください。
決め打ち二分探索することを考えます。整数 $0\le X \le M - 1$ を固定し、 $\displaystyle \min_{0\le x<N}\lbrace (Ax+B)\ \text{mod}\ M\rbrace \geq X$ になるか？という判定問題を考えます。この条件を同値変形していきます。
$
\begin{align*}
& \displaystyle \min_{0\le x<N}\lbrace (Ax+B)\ \text{mod}\ M\rbrace \geq X\\
\iff &\forall\ x\in [0,N-1],\ \lbrace (Ax+B)\ \text{mod}\ M\rbrace \geq X\\
\iff &\forall\ x\in [0,N-1],\ Ax+B-M\left\lfloor\frac{Ax+B}{M}\right\rfloor \geq X\\
\iff &\forall\ x\in [0,N-1],\ \left\lfloor\frac{Ax+B}{M}\right\rfloor \le \frac{Ax+B-X}M \\
\iff &\forall\ x\in [0,N-1],\ \left\lfloor\frac{Ax+B}{M}\right\rfloor \le \left\lfloor\frac{Ax+B-X}M\right\rfloor \\
\iff & \sum_{x=0}^{N-1}\left(\left\lfloor\frac{Ax+B-X}M\right\rfloor - \left\lfloor\frac{Ax+B}{M}\right\rfloor\right)=0\\
\iff & f(N,M,A,B)=f(N,M,A,B-X)
\end{align*}
$
ここで、後半の式変形では $0\le X\le M - 1$ より $\displaystyle \left\lfloor\frac{Ax+B-X}{M}\right\rfloor - \left\lfloor\frac{Ax+B}M\right\rfloor$ の取りうる値は $0$ か $-1$ のみであることを利用しています。
同値変形の最後の式が成立するかの判定は floor sum の $2$ 回の計算で $O(\log M)$ で可能です。したがって、判定問題が $O(\log M)$ で解け、二分探索で $O(\log M)$ 回判定が必要なので全体として $\displaystyle O\left(\left(\log M\right)^2\right)$ で解くことができます。
def floor_sum(n, m, a, b):
a1, a2 = a // m, a % m
s = n * (n - 1) // 2 * a1
b1, b2 = b // m, b % m
if a2 == 0:
return s + b1 * n
k = (a2 * (n - 1) + b2) // m
return s + n * (k + b1) - floor_sum(k, a2, m, m + a2 - b2 - 1)
for _ in range(int(input())):
n, m, a, b = map(int, input().split())
ok, ng = 0, m
fl = floor_sum(n, m, a, b)
while ng - ok != 1:
vs = (ok + ng) >> 1
if floor_sum(n, m, a, b - vs) == fl:
ok = vs
else:
ng = vs
print(ok)
余談ですが、この問題は $O(\log M)$ で解くこともできます。興味がある方は [Library Checker] Min of Mod of Linear - maspyのHP をご覧ください。
一般化
実は、 $\displaystyle f_{p,q}(N,M,A,B)=\sum_{i=0}^{N-1}i^p\left\lfloor\frac{Ai+B}{M}\right\rfloor^q$ という形の一般化 floor sum も同じ導出方法で解くことができます。
$\displaystyle S_k(N)=\sum_{i=0}^{N-1}i^k=\sum_{j=0}^{k+1}s_{k}(j)N^j$ として $S_k$ と $s_{k}$ を定義します。一般に、この $s_{k}(j)$ はファウルハーバーの公式などから求めることができます。(低次の場合は手計算の方が楽です)
$q=0$ なら $\displaystyle f_{p,0}(N,M,A,B)=\sum_{i=0}^{N-1}i^p=S_p(N)$ となるので、以降は $q>0$ の場合を考えます。
まずは普通の floor sum と同じように式変形することで $0\le A,B < M$ の場合に帰着できます。
また、$\displaystyle d_k=\left\lceil\frac{kM-B}A \right\rceil$ に対し、 $\displaystyle \left\lfloor\frac{Ai+B}{M}\right\rfloor=k \iff d_k\le i < d_{k+1}$ が成立する（ただし $d_{K+1}=N$）ので、 $\displaystyle K=\left\lfloor\frac{A(N-1)+B}{M}\right\rfloor$ として
$
\begin{align*}
&\phantom{=} f_{p,q}(N,M,A,B)\\
&= \sum_{i=0}^{N-1}i^p\left\lfloor\frac{Ai+B}{M}\right\rfloor^q\\
&=\sum_{k=0}^K \sum_{i=d_k}^{d_{k+1}-1}i^p\left\lfloor\frac{Ai+B}{M}\right\rfloor^q\\
&=\sum_{k=0}^K k^q\sum_{i=d_k}^{d_{k+1}-1}i^p\\
&=\sum_{k=0}^K k^q\left(S_p(d_{k+1})-S_p(d_k)\right)\\
&=K^qS_p(d_{K+1})-\sum_{k=1}^K\left(k^q-(k-1)^q \right)S_p(d_k)\\
&=K^qS_p(N)+\sum_{k=1}^K\sum_{i=0}^{q-1}\binom{q}{i}(-1)^{q-i}k^iS_p(d_k)\\
&=K^qS_p(N)+\sum_{i=0}^{q-1}\sum_{k=1}^K\binom{q}{i}(-1)^{q-i}k^i\sum_{j=0}^{p+1}s_p(j)d_k^j \\
&=K^qS_p(N)+\sum_{i=0}^{q-1}\sum_{j=0}^{p+1}\binom{q}{i}(-1)^{q-i}s_p(j)\sum_{k=1}^Kk^i\left\lfloor\frac{kM+A-B-1}A \right\rfloor^j \\
&=K^qS_p(N)+\sum_{i=0}^{q-1}\sum_{j=0}^{p+1}\binom{q}{i}(-1)^{q-i}s_p(j)\sum_{k=0}^{K-1}(k+1)^i\left\lfloor\frac{(k+1)M+A-B-1}A \right\rfloor^j \\
&=K^qS_p(N)+\sum_{i=0}^{q-1}\sum_{j=0}^{p+1}\binom{q}{i}(-1)^{q-i}s_p(j)\sum_{k=0}^{K-1}\sum_{l=0}^i\binom{i}{l}k^l \left\lfloor\frac{kM+M+A-B-1}A \right\rfloor^j \\
&=K^qS_p(N)+\sum_{i=0}^{q-1}\sum_{j=0}^{p+1}\sum_{l=0}^i\binom{q}{i}\binom{i}{l}(-1)^{q-i}s_p(j)\sum_{k=0}^{K-1}k^l \left\lfloor\frac{kM+M+A-B-1}A \right\rfloor^j \\
&=K^qS_p(N)+\sum_{i=0}^{q-1}\sum_{j=0}^{p+1}\sum_{l=0}^i\binom{q}{i}\binom{i}{l}(-1)^{q-i}s_p(j)f_{l,j}(K,A,M,M+A-B-1) \\
\end{align*}
$
となります。 $0\le l\le i\le q-1,0\le j\le p+1$ の範囲では $0\le l+j\le p+q$ となるので、 $0\le i+j\le p+q$ の範囲で各 $(N,M,A,B)$ のペアに対し $f_{i,j}(N,M,A,B)$ の値を全て保持しつつ再帰していくことで一般化 floor sum を解くことができます。
また、 $\displaystyle \sum_{i=0}^{q-1}\sum_{j=0}^{p+1}\sum_{l=0}^i\binom{q}{i}\binom{i}{l}(-1)^{q-i}s_p(j)f_{l,j}$ のテーブルは $N,M,A,B$ によらず一定なので、この部分を前計算することで計算量は $O((p+q)^4\log M)$ となります。
例えば $p+q\le 2$ の場合は以下のようになります。式の可視性を上げるために $B'=M-B+A-1$ としています。
$\displaystyle f_{0,1}(N,M,A,B)=NK-f_{0,1}(K,A,M,B')$
$\displaystyle f_{1,1}(N,M,A,B)=\frac{N(N-1)}2K+\frac12f_{0,1}(K,A,M,B')-\frac12f_{0,2}(K,A,M,B')$
$\displaystyle f_{0,2}(N,M,A,B)=NK^2-f_{0,1}(K,A,M,B')-2f_{1,1}(K,A,M,B')$
問題例
一般化 floor sum が使える問題を載せておきます。
yukicoder contest 394F - Inversion Number of Mod of Linear
一般化 floor sum が想定解です。
ARC182E - Sum of Min of Mod of Linear
想定解では一般化 floor sum を使用しませんが、一般化 floor sum を用いることで想定解中の必要な考察を飛ばすことができます。
以上が floor sum アルゴリズムとその一般化です。読んでいただきありがとうございました。