総和 $$ 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{aligned} 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{aligned} $$ と変形することができます。
また、$$ \begin{aligned} \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{aligned} $$ が成り立ちます。ここで、$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 につけていたタイトルです。
類問
- Codechef: https://www.codechef.com/problems/WINDOW/
- Rosecode: RoseCode: Problem #444 - A big sum 2
- LibreOJ: 类欧几里得算法 - 题目 - LibreOJ
- LibreOJ: 万能欧几里得 - 题目 - LibreOJ