問題を多項式の処理に帰着します。

箱 \(i\) に \(c_i\) 個のボールが入っている状態を多項式 \(x_1^{c_1}\cdots x_N^{c_N}\) で表します。操作を\(1\) 回行うことは \(\displaystyle\frac{1}{N}(x_1 + x_2 + \cdots + x_{N})\) をかけることに相当します。係数の \(\displaystyle\frac{1}{N}\) は確率の重みを表します。 各状態の スコア は、すべての \(x_i\) について偏微分を取り \(1\) を代入することで得られます。

\(y=x_1 + x_2 + \cdots + x_{N}\) とすると、 \(K\) 回の操作の後の状態は \(f=\displaystyle\frac{1}{N^K}\left(\displaystyle\prod_{i=1}^{N} x_i^{a_i}\right)\cdot y^K\) と表せます。 \(g=\displaystyle\frac{\partial^N}{\partial x_1\cdots\partial x_N} f\) を求めるのが目標です。ところで \(f\) の形の多項式について \(\displaystyle\frac{\partial}{\partial x_i}f=\left(\displaystyle\frac{a_i}{x_i}+\frac{\partial}{\partial y}\right) f\) が成立します。

これを用いて展開すると \(\displaystyle\frac{\partial^N}{\partial x_1\cdots\partial x_N} =\left(\displaystyle\frac{a_1}{x_1}+\frac{\partial}{\partial y}\right)\cdots \left(\displaystyle\frac{a_N}{x_N}+\frac{\partial}{\partial y}\right)\) のようになり、計算できます。

具体的には、この右辺の係数は \((a_1+X)\cdots(a_N+X)\) と同じになりますが、これは \(O(N^2)\) 、あるいは順番を工夫して FFT を使うと \(O(N(\log N)^2)\) で計算できます。