😊
floor_sum()

2024/09/23に公開
 イントロダクションいきなりですが、次の値を求めてください。

f(n, m, a, b) = \sum_{i=0}^{n-1} \Big\lfloor \frac{ai+b}{m} \Big\rfloor
O(n) で計算するアルゴリズムを思い浮かべた方も多いのではないでしょうか？　以下では、これをO\big(\log \max(a, m)\big)で計算するアルゴリズムを紹介します。

 アルゴリズム編!本節は参考記事に大いに影響を受けています。そちらも参照してください。

 Step 1まず、aとbを次のように展開します。

\begin{equation*}
    \begin{cases}
        a = d_a m + r_a \\
        b = d_b m + r_b \\
    \end{cases}
    ,\quad
    0 \le r_a, r_b \lt |m|
\end{equation*}
すると、f は次のようにかけます。

\begin{align*}
    f(n, m, a, b)
    &= \sum_{i=0}^{n-1}\Big\lfloor \frac{r_a i + r_b}{m} + d_a i + d_b \Big\rfloor \\
    &= \frac{n(n-1)}{2} d_a + n d_b + f(n, m, r_a, r_b)
\end{align*}
0 \le r_a, r_b \lt |m|より次が成り立ちます。

\begin{align*}
    f(n, m, r_a, r_b) &= (0 + \cdots + 0) + (1 + \cdots + 1) + \cdots + (k + \cdots + k) \\
    0 \le k &\coloneqq \Big\lfloor \frac{r_a (n-1) + r_b}{m} \Big\rfloor \le n - 1
\end{align*}

 Step 2: 主客転倒次の事実が成り立ちます。
0 \le i \le kを満たす任意のiについて、\lfloor (r_a j + r_b) / m \rfloor = iを満足するjが1つ以上存在する。
したがって、主客転倒により計算量の削減が見込めます。つまり、iの登場回数をO(1)で数えることができれば、O(k)で計算が終わります。そのためには、\lfloor (r_a j + r_b) / m \rfloor = iを満足する最大または最小のjを見つけることができればよいです。

\begin{equation*}
    i \le \frac{r_a j + r_b}{m} \lt i + 1
    \quad\Rightarrow\quad \frac{m i - r_b}{r_a} \le j_i
\end{equation*}
したがって、

\begin{align*}
    f(n, m, r_a, r_b)
    &= 0 \cdot (j_1 - j_0) + 1 \cdot (j_2 - j_1) + \cdots + k \cdot (n - j_{k}) \\
    &= kn - \sum_{i=1}^k j_i \\
    &= kn - \sum_{i=1}^k \Big\lceil \frac{mi - r_b}{r_a} \Big\rceil \\
\end{align*}

 Step 3: 再帰化次の事実が成り立ちます。
\lfloor x \rfloor = -\lceil -x \rceil
これを利用して、fを再起関数に落とし込みます。

\begin{align*}
    f(n, m, r_a, r_b)
    &= kn + \sum_{i=1}^k \Big\lfloor -\frac{mi - r_b}{r_a} \Big\rfloor \\
    &= kn + \sum_{i=0}^{k-1} \Big\lfloor -\frac{m (k - i) - r_b}{r_a} \Big\rfloor \\
    &= kn + \sum_{i=0}^{k-1} \Big\lfloor \frac{m i - m k + r_b}{r_a} \Big\rfloor \\
    &= kn + f(n, r_a, m, -m k + r_b)
\end{align*}
mとr_aが交換されていることに注意すると、計算量がユークリッドの互除法と同じになることが分かります。
ところで、r_a = 0のとき上式は破綻しています。f(n, m, 0, r_b) = n d_bより、r_a = 0を再帰の終了条件にすると回避できます。

 実装編実装上のポイントだけ簡単に解説します。
除算にdiv_euclid()やrem_euclid()を使用しています。通常の除算では余りが負になることがありますが、これらのメソッドでは常に正になることが保証されます。つまり、「切り捨て除算」を実行します。これによって、床関数の引数が負の場合でも正しく動作します。

mがゼロのとき、ゼロ除算になりpanicしてしまいます。これを防ぐためにstd::num::NonZeroI64を利用しても良いです。同様に、nは非負なので、u64を採用しても良いです。こうした型チェックの恩恵を受けるためには、関数のラッパーを用意すると良いです。
/// Calculate $\sum_{i=0}^{n-1} \lfloor \frac{ai+b}{m} \rfloor$ with $\O(\log a)$ time complexity.
fn floor_sum(n: i64, m: i64, a: i64, b: i64) -> i64 {
    assert_ne!(m, 0);
    assert!(n >= 0);

    let (div_a, rem_a) = (a.div_euclid(m), a.rem_euclid(m));
    let (div_b, rem_b) = (b.div_euclid(m), b.rem_euclid(m));

    let mut res = n * (n - 1) / 2 * div_a + n * div_b;
    if rem_a == 0 {
        return res;
    }

    let k = (rem_a * n - rem_a + rem_b).div_euclid(m);
    res += n * k;
    res += floor_sum(k, rem_a, m, rem_b - m * k);

    res
}

 Verify
 参考
イントロダクション

アルゴリズム編

Step 1

Step 2: 主客転倒

Step 3: 再帰化

実装編

Verify

参考

Discussion
