AGC019E (Shuffle and Swap) をたたみ込まないで解く

問題のリンク

解き方の概要

$A$ と $B$ ともに $1$ である添字の数を $X$ 、$A$ が $1$ 、 $B$ が $0$ である添字の数を $Y$ とします。($X+Y=k$)

$A$ の $0$ が何回を移動するかを考えます。 $Y$ 個の $0$ について、$i$ 個目の $0$ が $ a_{i} $ 回移動するとします。 $1$ と $1$ の無駄な交換が $X+Y- \sum_{i=1}^{Y} a_{i} $ 回あり、これを含めると交換の順序を求める式は以下のようになります。 $$ k! \prod_{i=1}^{Y}\frac{1}{a_{i}!} $$

$0$ の行き先の組み合わせは $Y!$ 通りあり、$ a_{i} $ を決めた後、 $0$ がどの場所を経由して移動するかと、 $1$ と $1$ の無駄な交換の決め方は $X!$ 通りあるので、 $\sum_{i=1}^{Y} a_{i} \leq X+Y$ を満たす、要素が正の、長さ $Y$ の全ての整数列 $a$ について$Y!X!k! \prod_{i=1}^{Y}\frac{1}{a_{i}!}$ を求めれば良いです。

ここで、 $f(x)=\sum_{i=1}^{\infty}\frac{x^{i}}{i!}$ とすると$(f(x))^{Y}$ の $x^{Y}$ から $x^{X+Y}$ の係数の和に $Y!X!k!$ をかけたものが答えです。

公式解説ではこれを繰り返し2乗法とたたみ込みを使うことで $O(n(\log n)^2)$ で計算してます。これをたたみ込まずに計算します。

$f(x)=(e^x -1)$ であることを利用すると以下のようになります。$$(f(x))^Y=\sum_{i=0}^{Y} \ _{Y} C_{i} \ e^{ix}(-1)^{Y-i}=\sum_{i=0}^{Y}\ _{Y} C_{i} (-1)^{Y-i}\sum_{j=0}^{\infty}\frac{(i)^{j}x^{j}}{j!}$$

係数が累乗の形と階乗と階乗の逆数の積の和で表すことができました。

これらの計算は $O(n^2)$ ですることができます。$n\leq 10^4$ ですが、$\frac{1}{4}$ の定数倍がつくので、実行時間制限に十分間に合います。

実装例