特に Advent Calendar は関係ない 12 月 13 日の記事です。
これの解説です。
atcoder.jp
以下の二つを参考にしながら考えたことをまとめます。
qiita.com
数式が多めです。もしかして MathJax から KaTeX に移行するべきですか？
（追記 2021/01/22：\(\KaTeX\) に移行しました）
問題概要
一次関数 \(f(x) = ax+b\) の出力を \(m\) で切り捨て除算した値を \(x\in[0, n)\) の整数について足したものを求めてね。
すなわち、以下で定義される \(g(n, m, a, b)\) を出力してね。
\[
g(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{f(i)}{m}\right\rfloor.
\]
重要な観察
\(g(n, m, a-m, b)\) と \(g(n, m, a, b-m)\) について調べてみます。
\[
\begin{aligned}
g(n, m, a-m, b) &= \sum_{i=0}^{n-1} \left\lfloor\frac{(a-m)\cdot i+b}{m}\right\rfloor \\
&= \sum_{i=0}^{n-1} \left\lfloor\frac{ai+b-im}{m}\right\rfloor \\
&= \sum_{i=0}^{n-1} \left(\left\lfloor\frac{ai+b}{m}\right\rfloor -i\right) \\
&= g(n, m, a, b) - \frac{n(n-1)}{2}.
\end{aligned}
\]
\[
\begin{aligned}
g(n, m, a, b-m) &= \sum_{i=0}^{n-1} \left\lfloor\frac{ai+(b-m)}{m}\right\rfloor \\
&= \sum_{i=0}^{n-1} \left(\left\lfloor\frac{ai+b}{m}\right\rfloor -1\right) \\
&= g(n, m, a, b) - n.
\end{aligned}
\]
さて、これによって以下のこともわかります（所望の回数だけ上記を適用する）。
\[
\begin{aligned}
g(n, m, a\bmod m, b\bmod m) &= g(n, m, a, b) - \left\lfloor\frac{a}{m}\right\rfloor \cdot \frac{n(n-1)}{2} - \left\lfloor\frac{b}{m}\right\rfloor \cdot n.
\end{aligned}
\]
すなわち、\(0\le a\lt m\) および \(0\le b\lt m\) に帰着できます。
あとは、これについて考えていきます。
考察
方針を先に書くと、小さいケースに帰着させて再帰的に解くことを目指します。
ベースケースは、\(g(0, m, a, b) = 0\) です。
いわゆる contribution technique とか 主客転倒 とか呼ばれているテクニックで総和をバラしていきます。ある値を見たときに、それが何回総和に寄与するかを数えるやつです。
\(\gdef\Y{y_{\text{max}}}\)
\(\Y = \lfloor f(n-1)/m\rfloor = \lfloor (a(n-1)+b)/m\rfloor\) とおきます。
（追記 2020/12/29：あるいは \(\Y = \lfloor f(n)/m\rfloor\) とおきます。両方試してみましょう）
\[
\begin{aligned}
g(n, m, a, b) &= \sum_{i=0}^{n-1} \left\lfloor\frac{f(i)}{m}\right\rfloor \\
&= \sum_{j=1}^{\Y} \sum_{\substack{0\le i\lt n\\ \lfloor f(i)/m \rfloor=j}} j. \\
\end{aligned}
\]
ここで、\(k_j\) を \(\lfloor f(i)/m\rfloor = j\) なる最小の \(i\) とおく*1と、次のように書けます。
（追記 2021/01/10：行間を少しうめました）
\[
\begin{aligned}
g(n, m, a, b) &= \sum_{j=1}^{\Y} \sum_{\substack{0\le i\lt n\\ \lfloor f(i)/m \rfloor=j}} j \\
&= \sum_{j=1}^{\Y-1} j\cdot(k_{j+1}-k_j) + \Y\cdot(n-k_{\Y}) \\
&= \sum_{j=1}^{\Y-1} j\cdot k_{j+1} - \sum_{j=1}^{\Y-1} j\cdot k_j + \Y\cdot n - \Y\cdot k_{\Y} \\
&= \sum_{j=2}^{\Y} \,(j-1)\cdot k_j - \sum_{j=1}^{\Y-1} j\cdot k_j + \Y\cdot n - \Y\cdot k_{\Y} \\
&= \sum_{j=2}^{\Y} \,(j-1)\cdot k_j - \sum_{j=1}^{\Y} j\cdot k_j + \Y\cdot n \\
&= \sum_{j=2}^{\Y} \,(j-1)\cdot k_j - 1\cdot k_1 - \sum_{j=2}^{\Y} j\cdot k_j + \Y\cdot n \\
&= - 1\cdot k_1 - \sum_{j=2}^{\Y} - k_j + \Y\cdot n \\
&= \sum_{j=1}^{\Y} -k_j + n\cdot\Y.
\end{aligned}
\]
さて、\(k_j\) を求めます。
定義から、以下が成り立ちます。
\[
\left\lfloor\frac{f(k_j-1)}{m}\right\rfloor \lt j = \left\lfloor\frac{f(k_j)}{m}\right\rfloor \le \frac{f(k_j)}{m}.
\]
ここで \(j \le f(k_j-1)/m\) と仮定すると、\(j\) は整数なので \(\lfloor f(k_j-1)/m\rfloor \ge j\) となり、\(\lfloor f(k_j-1)/m\rfloor \lt j\) に矛盾します。
よって、以下を得ます。
（追記 2021/01/10：2 行目から 3 行目は、左と右の不等式をそれぞれ \(k_j\) について変形してまとめます。）
\[
\begin{aligned}
\frac{f(k_j-1)}{m} \lt j \le \frac{f(k_j)}{m}; \\
\frac{a(k_j-1)+b}{m} \lt j \le \frac{ak_j+b}{m}; \\
\frac{mj-b}{a} \le k_j \lt \frac{mj-b}{a}+1; \\
k_j = \left\lceil \frac{mj-b}{a} \right\rceil.
\end{aligned}
\]
あとは、これを元の式に入れてがんばります。
\(\lceil x\rceil = -\lfloor -x\rfloor\) を使ったり、添字を入れ替えたりします*2。
\[
\begin{aligned}
g(n, m, a, b) &= \sum_{j=1}^{\Y} -k_j + n\cdot\Y \\
&= n\cdot \Y + \sum_{j=1}^{\Y} - \left\lceil \frac{mj-b}{a} \right\rceil \\
&= n\cdot \Y + \sum_{j=1}^{\Y} + \left\lfloor \frac{-mj+b}{a} \right\rfloor \\
&= n\cdot \Y + \sum_{j=0}^{\Y-1} + \left\lfloor \frac{-m(\Y-j)+b}{a} \right\rfloor \\
&= n\cdot \Y + \sum_{j=0}^{\Y-1} + \left\lfloor \frac{mj+b-m\cdot\Y}{a} \right\rfloor \\
&= n\cdot \Y + g(\Y, a, m, b-m\cdot\Y).
\end{aligned}
\]
より小さいケースに帰着できました。
念のため、小さくなっていることを確認しておきます。\(0\le a\lt m\)、\(0\le b\lt m\)、\(1\le n\) に注意です。
（追記 2021/01/10：ここの議論は、\(\Y = \lfloor f(n)/m\rfloor\) のときは少し変わると思います。）
\[
\begin{aligned}
n - \Y &= n-\left\lfloor\frac{a(n-1)+b}{m}\right\rfloor \\
&\ge n - \frac{a(n-1)+b}{m} \\
&= \frac{mn-an+a-b}{m} \\
&\gt \frac{mn-an+a-m}{m} \\
&= \frac{(m-a)(n-1)}{m} \ge 0; \\
n &\gt \Y.
\end{aligned}
\]
ここで終わってもいいのですが、もう一工夫です。
上の重要な観察のところで示したように、第 4 引数から第 2 引数を好きな回数だけ引いた値も計算できるので、それをしてみます。
具体的には、\(-(n-1)\) 回だけ引きます。
\[
g(\Y, a, m, a(n-1)+b-m\cdot\Y) = g(\Y, a, m, b-m\cdot\Y) + (n-1)\cdot\Y.
\]
すなわち、次の式を得ます。\(\Y = \lfloor f(n-1)/m\rfloor\) を思い出しましょう。
（追記 2021/01/10：\(x\bmod y = x - \lfloor x/y\rfloor\cdot y\) も使います。）
\[
\begin{aligned}
g(\Y, a, m, b-m\cdot\Y) &= g(\Y, a, m, a(n-1)+b-m\cdot\Y) - (n-1)\cdot\Y \\
&= g(\Y, a, m, f(n-1)-m\cdot\lfloor f(n-1)/m\rfloor) - (n-1) \cdot \Y \\
&= g(\Y, a, m, f(n-1)\bmod m) -(n-1)\cdot\Y.
\end{aligned}
\]
これを元の式に戻してゴールです。
\[
\begin{aligned}
g(n, m, a, b) &= n\cdot \Y + g(\Y, a, m, b-m\cdot\Y) \\
&= n\cdot \Y + g(\Y, a, m, f(n-1)\bmod m) -(n-1)\cdot\Y \\
&= \lfloor f(n-1)/m\rfloor + g(\lfloor f(n-1)/m\rfloor, a, m, f(n-1)\bmod m).
\end{aligned}
\]
\(\lfloor f(n-1)/m\rfloor\) と \(f(n-1) \bmod m\) は、コンパイラを信じれば除算命令 1 回でやってくれると思います。
実装
fn floor_sum(n: i64, m: i64, mut a: i64, mut b: i64) -> i64 {
let mut res = 0;
if a >= m {
res += n * (n - 1) / 2 * (a / m);
a %= m;
};
if b >= m {
res += n * (b / m);
b %= m;
};
let last = a * (n - 1) + b;
if last >= m {
let last_div = last / m;
let last_mod = last % m;
res += last_div + floor_sum(last_div, a, m, last_mod);
}
res
}
適当なループをして非再帰に書き直すことは可能だと思いますが、このくらいならコンパイラに任せた方がいいかもしれません。
ループにしても大して速くならず、無駄に可読性を落とすことが予想されます。
（追記：2020/12/16）
よく考えると、last のところは a * n + b にできて、そうすると last_div を足す必要がなくなって、次のようにできます。
\(\Y\) や上の式変形ででてきた \(n-1\) を \(n\) に置き換えてよいので。
fn floor_sum(n: i64, m: i64, mut a: i64, mut b: i64) -> i64 {
let mut res = 0;
if a >= m {
res += n * (n - 1) / 2 * (a / m);
a %= m;
};
if b >= m {
res += n * (b / m);
b %= m;
};
let last = a * n + b;
if last >= m {
res += floor_sum(last / m, a, m, last % m);
}
res
}
なんですが、速くなってくれませんでした（以下の提出コードは諸事情で C++ です）。
提出し直したら速くなったりしてわかりません。ジャッジの気分屋さん。
停止性の証明だけ書き直します。（訂正 2021/01/10：やめました）
冷静に考えると、\(n-\Y\gt 0\) を示さなくても、計算量のところに書くように \(a\) と \(m\) の互除法より長くはループしないことがわかるので十分な気もします。
（追記 2021/01/10：負数に対応させる場合は、/ が切り捨て除算でない言語では注意が必要です。）
計算量
（追記：2020/12/14）
\(m, a\) のみに注目すると互除法のようになっていて、最悪ケースでは \(\log_{\varphi}(\min\{m, a\})\) 回程度の再帰呼び出しをします。
互除法が終わったときの \(\Y\) の値を考えると、\(a = 0\) かつ \(b\lt m\) なので、それ以上の呼び出しはしないことがわかります。
\(\varphi\) は黄金比です。
また前述のように、呼び出しごとに \(n\) が少なくとも \(\frac{(m-a)(n-1)}{m}\) 減ります。
互除法の最悪ケースでは \(a/m\) は \(1/\varphi\) くらいになっているので、ざっくり高々 \(\log_{\varphi}(n)\) 回程度再帰呼び出ししそうな気がします。
怪しそう。たすけて〜。
より詳細に見積もる方法を身につけたいです。変数がいっぱいあって大変。
結局、再帰呼び出しは \(\log_{\varphi}(\min\{m, a, n\})\) 回程度で抑えられて、各々は \(O(1)\) 時間なので、全体として \(O(\log(\min\{m, a, n\}))\) 時間だと思います。
big-O だと隠されちゃうけど、それはそれとしてループ回数を詳しく解析するのは楽しそうなので、できるようになりたいです。
\(n\) に関する評価のきもち
（追記 2021/01/10）
Fibonacci 数列を \(\{F_i\}\) とします。
\(a\) と \(m\) の互除法の最悪ケースでは、 \( (F_i, F_{i-1}) \dots, (F_2, F_1), (F_1, F_0)\) のように動きます。
このとき、\(j\) ステップ後の \(n\) は次のようになっているはずです（\(\pm 1\) とかは追々考える必要があります）。
\[
\left\lfloor \frac{F_{i-j}}{F_{i-j+1}} \dots \left\lfloor \frac{F_{i-1}}{F_i}\cdot n\right\rfloor\dots\right\rfloor
\approx \left\lfloor \frac{F_{i-j}}{F_i}\cdot n \right\rfloor.
\]
\(n\) が十分小さければ、\(j=i\) となる前にこれは \(0\) になるはずです。
\(F_i が \varphi^i\) くらいであることを考えると、\(\log_{\varphi}(n)\) ステップくらいで \(0\) になりそうです。
\(a\) と \(m\) については互除法の議論から \(\log_{\varphi}(\min\{a, m\})\) ステップくらいで \(0\) になるので、これらの \(\min\) のオーダーになりそうです。
おわり
お疲れさまでした。
うれしい
この高速化の PR が ACL にマージされました。
うれしい 2
yosupo さんが ACL Practice Contest の 解説 に載せてくれました。
