floor sum アルゴリズムとその一般化

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 で書くと以下のようになります。

 ```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)](https://judge.yosupo.jp/problem/min_of_mod_of_linear)

整数 $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)$ で解くことができます。

[実装例(Python3)](https://judge.yosupo.jp/submission/255749)
```Python3
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](https://maspypy.com/library-checker-min-of-mod-of-linear) をご覧ください。

# 一般化

実は、 $\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=0}^{K-1}\left((k+1)^q-k^q \right)S_p(d_{k+1})\\\\
&=K^qS_p(d_{K+1})-\sum_{k=0}^{K-1}\sum_{i=0}^{q-1}\binom qik^iS_p(d_{k+1})\\\\
&=K^qS_p(d_{K+1})-\sum_{k=0}^{K-1}\sum_{i=0}^{q-1}\binom qik^i\sum_{j=0}^{p+1}s_p(j)d_{k+1}^j\\\\
&=K^qS_p(d_{K+1})-\sum_{k=0}^{K-1}\sum_{i=0}^{q-1}\binom qik^i\sum_{j=0}^{p+1}s_p(j)\left\lceil\frac{(k+1)M-B}A \right\rceil^j\\\\
&=K^qS_p(d_{K+1})-\sum_{i=0}^{q-1}\binom qi\sum_{j=0}^{p+1}s_p(j)\sum_{k=0}^{K-1}k^i\left\lfloor\frac{(k+1)M-B+A-1}A \right\rfloor^j\\\\
&=K^qS_p(d_{K+1})-\sum_{i=0}^{q-1}\binom qi\sum_{j=0}^{p+1}s_p(j)f_{i,j}(K,A,M,M-B+A-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}\binom qi\sum_{j=0}^{p+1}s_p(j)f_{i,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](https://yukicoder.me/problems/no/2362)

一般化 floor sum が想定解です。

[ARC182E - Sum of Min of Mod of Linear](https://atcoder.jp/contests/arc182/tasks/arc182_e)

想定解では一般化 floor sum を使用しませんが、一般化 floor sum を用いることで想定解中の必要な考察を飛ばすことができます。

---

以上が floor sum アルゴリズムとその一般化です。誤植やより良い式変形などを指摘してくださった potato167 さん、p進大好きbot さん、かならい さん、noshi91 さん、ありがとうございます。
読んでいただきありがとうございました。