Binomial Sums

$(\sum_{i=0}^{n}\binom{n}{i}i^k) \bmod p$ を $O(k + \log(n))$ で計算する。条件: $n \ge k$, $p$: 素数, $p > k$

Hackerrank の Costly Graphs での bayleef さんの解法の導出です(Editorial には導出方法が載っていないため)。母関数を使うのですが、あまり見慣れない変形があります。

まず、数列 $\binom{n}{0}0^k, \binom{n}{1}1^k, \dots$ の母関数は、$$\begin{aligned} \left(x\frac{d}{dx}\right)^k (1+x)^n &= \sum_{i=0}^k \genfrac{\{}{\}}{0pt}{}{k}{i} x^i \frac{d^i}{dx^i} (1 + x)^n \\ &= \sum_{i=0}^{k} \genfrac{\{}{\}}{0pt}{}{k}{i}\binom{n}{i}i! \cdot x^i(1+x)^{n-i} \end{aligned}$$ となります。ここで、 $\genfrac{\{}{\}}{0pt}{}{n}{k}$ は第 2 種 Stirling 数を示します。

これは $n$ 次多項式として見れるので、$x = 1$ とすると、 $$\sum_{i=0}^{n}\binom{n}{i}i^k = \sum_{i=0}^{k} \genfrac{\{}{\}}{0pt}{}{k}{i}\binom{n}{i}i! 2^{n-i} $$ となり、これは FFT/NTT を用いて第 2 種 Stirling 数を計算すると $O(k \log(k) + \log(n))$ で求まります(想定解法)。

ただ、この式はまだ変形することができて、$$ \begin{aligned}\sum_{i=0}^{k} \genfrac{\{}{\}}{0pt}{}{k}{i}\binom{n}{i}i! 2^{n-i} &= 2^{n-k} \sum_{i=0}^{k} \genfrac{\{}{\}}{0pt}{}{k}{i}\binom{n}{i}i! 2^{k-i} \\ &= 2^{n-k} [x^k] \frac{\sum_{i=0}^{k} \genfrac{\{}{\}}{0pt}{}{k}{i}\binom{n}{i}i! \cdot x^i(1+x)^{k-i} }{1-x} \\ &= 2^{n-k} [x^k] \left(\frac{\sum_{i=0}^{k} \genfrac{\{}{\}}{0pt}{}{k}{i}\binom{n}{i}i! \cdot x^i(1+x)^{n-i} }{(1-x)} \cdot \frac{1}{(1+x)^{n-k}}\right) \\ &= 2^{n-k} [x^k] \left( \left(\sum_{i=0}^\infty \left(\sum_{j=0}^{i} \binom{n}{j}j^k\right)x^i \right) \cdot \left(\sum_{i=0}^\infty \binom{n-k -1+i}{i}(-1)^i x^i \right) \right) \ \end{aligned}$$ とすることができます。

よって、 $$\sum_{i=0}^{n}\binom{n}{i}i^k = 2^{n-k} \left( \sum_{i=0}^{k} \left( \sum_{j=0}^i \binom{n}{j}j^k \right) \cdot \binom{n-1-i}{k-i}(-1)^{k-i} \right)$$ として計算でき、これは線形篩を用いれば mod $p$ で $O(k + \log(n))$ で計算できます。

また、$i^k$ の部分が $k$ 次以下の多項式 $f$ であった場合でも、同様の式変形が行えるため、$$\sum_{i=0}^{n}\binom{n}{i} f(i) = 2^{n-k} \left( \sum_{i=0}^{k} \left( \sum_{j=0}^i \binom{n}{j} f(j) \right) \cdot \binom{n-1-i}{k-i}(-1)^{k-i} \right)$$ と計算できます。