Official

Ex - Dye Color Editorial by PCTprobability


数列 \(A=(A_1,A_2,\dots,A_N)\) に対して、操作回数の期待値を \(f(A)\) と置きます。また、\(A\) の全ての要素が等しいことを終端状態ということにします。

数列 \(A\) から \(1\) 回の操作で数列 \(A'\) になる確率を \(E_{A,A'}\) と置きます。

\(A\) が終端状態である場合は \(f(A)=0\)、そうでない場合は \(f(A) = 1 + \sum E_{A,A'} f(A')\) が成り立ちます。

また、\(B_{A,j} := (A_i = j\) を満たす \(i\) の個数 \()\) と数列 \(B\) を定義します。

ある定数 \(C\) が存在し、\( \sum_{i=1}^{N} g(B_{A,i}) = f(A) + C\) となる関数 \(g(x)\) が存在すれば本問題を解くことができます。

\(g\) が満たすべき式は終端状態でない数列 \(A\) に対する \( \sum_{i=1}^{N} g(B_{A,i}) =1+ \sum E_{A,A'} \sum_{i=1}^{N} g(B_{A',i})\) です。

もし、\(g\)\(1 \le i \le N\) に対して \(g(B_{A,i}) = \frac{1}{N} + \sum E_{A,A'} g(B_{A',i})\) を満たせば上記の式は必ず満たされます。

更に、今 \(B_i = j\) である時、\(1\) 回操作した後に \(B_i = k\) となる確率は \(j,k\) にのみ依存する形で表すことができます。

より、\(1\) 回操作した後に \(B_i = j\) から \(B_i = k\) となる確率を \(P_{j,k}\) と置くと、\(g\) の満たすべき条件は \(0 \le i < N\) に対して、\(g(i) = \frac{1}{N} + \sum_{j=0}^{N} P_{i,j} g(j)\) です。\(i=N\) に対しては、既に終端状態となっているため上記の式を満たす必要はありません。

更に、\(1\) 回の操作で \(B_i\)\(2\) 以上増えることはあり得ません。

なので、この連立方程式を行列で表すと以下のようになります。

\(\begin{pmatrix} P_{0,0}-1 & P_{0,1} & 0 & 0 & \dots & 0 & 0\\ P_{1,0} & P_{1,1}-1 & P_{1,2} & 0 & \dots & 0 & 0\\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots\\ P_{N-1,0} & P_{N-1,1} & P_{N-1,2} & P_{N-1,3} & \dots & P_{N-1,N-1}-1 & P_{N-1,N} \\ \end{pmatrix}\) \(\begin{pmatrix} g(0) \\ g(1) \\ \vdots \\ g(N) \\ \end{pmatrix}\) \(=\) \(\begin{pmatrix} -\frac{1}{N} \\ -\frac{1}{N} \\ \vdots \\ -\frac{1}{N} \\ \end{pmatrix}\)

\(g(0)=0\) とし、\(i=1,2,\dots,N-1\) の順番で \(i\) 行目の制約を利用することで \(\mathrm{O}(N^2)\)\(g(i)\) を求めることができます。

しかし、まだ課題が残っています。それは \(P_{i,j}\) を高速に求めることと、\(P_{i,i+1} \not \equiv 0 \bmod 998244353\) を証明することです。

\(P_{i,j}\) を高速に求めます。つまり、色 \(X\)\(i\) 個ある状況で、\(1\) 回の操作で \(j\) 個になる確率を求めます。

\(X\) のボールが初めに \(S\) 個選ばれたとします。その時、ボールの色を変更する時点で色 \(X\) が選ばれる確率は \(\displaystyle \frac{\displaystyle \sum_{k=0}^{N-i} \binom{N-i}{k} \times \binom{i}{S} \times \frac{k+S}{N}}{2^N}\) です。 \(k\) に依存する項、\(S\) に依存する項、\(k+S\) に依存する項、定数項に分けることができるため、畳み込みを使い \(\mathrm{O}(N\log N)\) で全ての \(S\) に対し上記の値を求めることができます。\(\frac{k+S}{N}\)\(\frac{k}{N} + \frac{S}{N}\) に分けるなどによって \(\mathrm{O}(N)\) で求めることも可能です。

ボールの色を変更する時点で色 \(X\) が選ばれない確率も同様にして \(\mathrm{O}(N\log N)\) で求めることができます。

より、\(P_{i,j}\)\(\mathrm{O}(N^2\log N)\) で全て求めることができました。

そして、上記の \(S = 0\) かつボールの色を変更する時点で色 \(X\) が選ばれる確率が \(P_{i,i+1}\) であり、その値は \(\displaystyle \frac{\displaystyle \sum_{k=0}^{N-i} \binom{N-i}{k} \times \frac{k}{N}}{2^N}\) です。式変形すると、この値は \(\displaystyle \frac{N-i}{N \times 2^{i+1}}\) となり、\(0 \le i < N\) に置いて \(\bmod 998244353\) で非零であることがわかりました。

よって、この問題を \(\mathrm{O}(N^2\log N)\)\(\mathrm{O}(N^2)\) で解くことができました。

posted:
last update: