sum of floor of linear の解説するよ

特に Advent Calendar は関係ない 12 月 13 日の記事です。

これの解説です。

atcoder.jp

以下の二つを参考にしながら考えたことをまとめます。

ACL の floor_sum のコード (特に x_max まわり) が何やってるかわからなくて自分で考え直してたら、x_max が必要なくなってついでに30%くらい高速なりました🤔https://t.co/uZ7GUND73d

— fura (@fura_2) 2020å¹´10月30æ—¥

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++ です）。 提出し直したら速くなったりしてわかりません。ジャッジの気分屋さん。

• 上のやつ: 70 ms
• 下のやつ: 75 ms

停止性の証明だけ書き直します。（訂正 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 の 解説 に載せてくれました。

*1:\(0\le b\lt m\) より、各 \(1\le j\le \Y\) について \(k_j\) は存在します。

*2:\(j\) の係数を非負にしたいので。