Scary Sum

総和 $$ S(a, b, k) := \sum_{x=0}^{b-1} \left\lfloor\frac{ax}{b}\right\rfloor^k $$ を $\bmod{p}$（十分に大きな素数）で $O(k \log k \log(\min\{a, b\}))$ で求めます。ただし、$\gcd(a, b) = 1$ とします。

求め方

仮定より $\gcd(a, b) = 1$ であるので、$a/b$ の部分を Stern–Brocot tree を用いて分解していくことができます。

例えば、非負整数 $a, b, c, d$ ($b, d > 0$) が $bc - ad = 1$ を満たしているとすると、$$ \begin{align*}
S(a + c, b + d, k) &= \sum_{x=0}^{b+d-1} \left\lfloor \frac{a+c}{b+d} \cdot x\right\rfloor ^k \\
&= \sum_{x=0}^{b-1} \left\lfloor\frac{ax}{b}\right\rfloor^k + \sum_{x=0}^{d-1} \left(\left\lfloor\frac{cx}{d}\right\rfloor + a\right)^k
\end{align*} $$ と変形することができます。

また、$$ \begin{align*}
\sum_{x=0}^{b-1} \left(\left\lfloor\frac{ax}{b}\right\rfloor + q\right)^k
&= k! \sum_{i=0}^k \frac{q^{k-i}}{(k-i)!} \cdot \frac{S(a,b,i)}{i!}, \\
\sum_{x=0}^{qb-1} \left\lfloor\frac{ax}{b}\right\rfloor^k
&= \sum_{j=0}^{q-1} \sum_{x=0}^{b-1} \left(\left\lfloor\frac{ax}{b}\right\rfloor + aj\right)^k \\
&= k! \sum_{i=0}^k \frac{a^{k-i} \cdot P_{k-i}(q-1)}{(k-i)!} \cdot \frac{S(a,b,i)}{i!}
\end{align*} $$ が成り立ちます。ここで、$P_k(n) := \sum_{x=0}^{n} x^k$ です。（2 つ目の式は、例えば、$\frac{7}{11} = \frac{1 + 3 \cdot 2}{2 + 3 \cdot 3}$ のように効率よく分解する場合に用います。）

上の 2 式は畳み込みの形になっているので、$P_i(q-1)$ の値を計算できているならば $O(k \log k)$ で計算可能ですが、$P_k(n)$ は Bernoulli 数 $B_i$ ($i = 0, \dots, k$) を用いて $$ P_k(n) = \sum_{x=0}^{n} x^k = k! \sum_{i=0}^{k} \frac{B_{i}}{i!} \cdot \frac{(n + 1)^{k+1-i}}{(k+1-i)!} $$ とかけ、Bernoulli 数 $B_i$ ($i = 0, \dots, k$) は、$$ \frac{x}{e^x - 1} \bmod x^{k+1} $$ の係数を Newton 法を用いて計算することで $O(k \log k)$ で求めることができるので、実際 $O(k \log k)$ で計算可能です。

分数の分解の回数については、愚直に行うと $O(b)$ 回かかりますが、効率よく行うと $O(\log({\min\{a, b\}}))$ 回で済ませることができます。

def gcd(a, b):
  while b:
    a, b = b, a % b
  return a

def decompose(a, b):
  assert(gcd(a, b) == 1)
  ln, ld = a // b, 1
  rn, rd = ln + 1, 1
  print("(%d/%d, %d/%d)" % (ln, ld, rn, rd))
  while (ln, ld) != (a, b) and (rn, rd) != (a, b):
    mn, md = (ln + rn), (ld + rd)
    if mn * b <= a * md:
      q = (a * ld - b * ln) // (b * rn - a * rd)
      ln, ld = ln + rn * q, ld + rd * q
      rn, rd = ln + rn, ld + rd
    else:
      q = (b * rn - a * rd) // (a * ld - b * ln)
      rn, rd = ln * q + rn, ld * q + rd
      ln, ld = rn + ln, rd + ld
    print("(%d/%d, %d/%d)" % (ln, ld, rn, rd))
  print(a, b)

decompose(3 ** 19, 5 ** 13)

よって、$S(a, b, k)$ は $\bmod p$ で $O(k \log k \log{(\min\{a, b\}}))$ で求めることができます。

補足

  • 実際には、$\gcd(a, b) = 1$ は不要です。
  • 総和の部分は $b-1$ ではなく、任意の整数 $n$ でも問題ないはず（場合分けが増える & 未実装）。
  • $\sum_{x=0}^{n} \left\lfloor\frac{ax + b}{c}\right\rfloor^k$ も、範囲をずらすことで、この問題に帰着できる（未実装）。
  • $\sum_{x=0}^{n} x^{k_1}\left\lfloor\frac{ax + b}{c}\right\rfloor^{k_2}$ は、2D FFT を使えば $O(k_1 k_2 \log{(k_1 k_2)} \log{(\min\{a, c\})})$ でできるはず（未確認）。
  • Scary Sum は Anton Lunyov さんが故 Mathalon の Problem 145 につけていたタイトルです。

類問