Ex - Dye Color Editorial
by 
potato167
線形
以下のように FPS(形式的冪級数) を定義します。
\(F(x) : \) 全てが同じ色である状態から、\(a\) 回操作した後に全てが同じ色になっている確率が \([x^{a}]F(x)\) であるような FPS
\(G(x) : \) 初期状態から \(a\) 回操作した後に全てが同じ色になっている確率が \([x^{a}]G(x)\) であるような FPS
\(H(x) :\) 初期状態から \(a\) 回操作した後に全てが同じ色かつ、途中で全てが同じ色になることがなかった確率が \([x^{a}]H(x)\) であるような FPS
この解説で出てくる FPS は全て有理関数(\(\frac{多項式}{多項式}\))です。
また、\(\mathrm{freq}(i)\) を \(A\) の中にある \(i\) の個数とします。
この問題の答えは、\(H'(1)\) と表せます。 \(H, F, G\) に対して、以下の関係式が成り立ちます。
\[H(x)F(x) = G(x)\]
\[H(x) = \dfrac{G(x)}{F(x)}\]
\[H'(x) = \dfrac{F(x)G'(x) - G(x)F'(x)}{F(x)^{2}}\]
\(F(x), G(x)\) がどのような形なのかを考えます。
ここで、任意の非負整数 \(a\) に対して、以下が成り立つように、 \(i = 0, 1, \dots , N\) に対して、\(f_{i}(x)\) を定義します。
- \(a\) 回操作して、ボール \(1, 2, \dots i\) が操作によって色 \(1\) になり、ボール \(i + 1, \dots , N\) が一度も塗られない確率が \([x^{a}]f_{i}(x)\) と等しい。
問題の操作は、ランダムに \((1, 2, \dots , N)\) の順列 \(P\) と、\(0\) 以上 \(2^{N}\) 未満の整数 \(X\) を選び、\(2\) 進数表記で \(X\) の \((i - 1)\) 桁目が \(1\) ならば、ボール \(i\) を \(P_{i}\) 色に塗ることと同値です。
この操作を後ろから考えることで、 \(f_{i}(x)\) に対して、以下が成り立つことが示せます。
\[f_{i}(x) = i!\prod_{k = 0}^{i}\dfrac{1}{1 - \frac{2^{k}x}{2^{N}}} \prod_{k = 0}^{i - 1}\dfrac{2^{k}x}{2^{N}N}\]
この \(f_{i}(x)\) は \(f_{N}(x)\) のみが分母に \((1 - x)\) を持つことを後に使います。
この \(f_{i}(x)\) を用いて、\(F(x), G(x)\) は以下のように表せます。
\[F(x) = \left(\sum_{i = 0}^{N - 1}\left( \begin{matrix} N\\ i\\ \end{matrix} \right)f_{i}(x)\right)+ N\cdot f_{N}(x)\]
\[G(x) = \sum_{n = 1}^{N}\left(\sum_{i = N - \mathrm{freq}(n)}^{N}\left( \begin{matrix} \mathrm{freq}(n)\\ N - i\\ \end{matrix} \right)f_{i}(x)\right)\]
\(F, G\) ともに \(f_{N}(x)\) にかかる係数が \(N\) であることから、以下が成り立つような、\((1 - x)\) を分母に含まない多項式 \(A(x), B(x), C(x)\) が存在します。
\[F(x) = A(x) + \dfrac{C(x)}{1 - x}\]
\[G(x) = B(x) + \dfrac{C(x)}{1 - x}\]
この \(A, B, C\) を用いて、\(H'(x)\) を表すと、以下のようになります。
\[H'(x) = \dfrac{F(x)G'(x) - G(x)F'(x)}{F(x)^{2}} = \dfrac{(1 - x)(hoge\dots) + (A(x) - B(x))C(x)}{(A(x)(1 - x)+C(x))^{2}}\]
よって、以下が成り立ちます。
\[H'(1) = \dfrac{A(1) - B(1)}{C(1)}\]
上記の式に基づいて、\(A(1), B(1), C(1)\) の値を求めれば良いです。\(M=998244353\) として、時間計算量は \(O(N\log{M})\) です。
この計算量は、\(i = 0 , 1, \cdots N - 1\) に対して \(1 - \dfrac{2^{i}}{2^{N}}\) の逆数を求める部分がボトルネックとなっていました。
一般に、\(M\) が素数である状況において、任意の要素が \(\bmod{M}\) で \(0\) でない数列 \(V = (V_{1}, V_{2}, \dots, V_{|V|})\) の全ての要素に対して逆数を求めるのは \(O(|V| + \log{M})\) で計算できます。これは \(\prod V\) の逆数に対して、\(V\) の後ろの要素から順に積をとることで、prefix 総積の逆数が求まり、その逆数に対し、一個前の prefix の積をかけることである要素に対する逆数の値がもとまるからです。
これと同様の方法で逆数を求めることで、全体の計算量が \(O(N + \log{M})\) となります。