😊

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

まず、abを次のように展開します。

\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*}

mr_aが交換されていることに注意すると、計算量がユークリッドの互除法と同じになることが分かります。

ところで、r_a = 0のとき上式は破綻しています。f(n, m, 0, r_b) = n d_bより、r_a = 0を再起の終了条件にすると回避できます。

実装編

実装上のポイントだけ簡単に解説します。

  1. 除算にdiv_euclid()rem_euclid()を使用しています。通常の除算では余りが負になることがありますが、これらのメソッドでは常に正になることが保証されます。つまり、「切り捨て除算」を実行します。これによって、床関数の引数が負の場合でも正しく動作します。
  2. 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;
    }

    // contribution technique
    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

参考

Discussion