G - Socks 3 Editorial
by
Ebishu0309
期待値 DP の高速化
同じ色の靴下を合計で \(2\) 枚取り出した時点で操作が終了するので、各色の靴下をそれぞれ合計 \(1\) 枚まで取り出している状態を考えます。
まず、次のような DP を考えます。
- \( \mathrm{dp}_{S}\) : 取り出した靴下の色の集合が \(S\) の状態から、操作が終了するまでに靴下を取り出す回数の期待値
この DP は以下の漸化式で計算できますが、これでは間に合いません。
\[ \mathrm{dp}_S = 1 + \sum_{i \notin S} \frac{A_i}{(\sum A)-|S|} \mathrm{dp}_{S\cup\{i\}} \]
そこで、\(\mathrm{dp}_S\) の値をまとめて計算することを考えます。\(|S|\) が同じならば分母の \((\sum A)-|S|\) が同じであることと、状態 \(S\) から状態 \(S\cup\{i\}\) に遷移する際に \(A_i\) を乗じていることを考えると、以下のようにまとめて計算することができます。
\[ f_m = m!\sum_{|S|=m}\prod_{i\in S}A_i\,\mathrm{dp}_S\]
として、
\[ f_m = m!\sum_{|S|=m}\prod_{i\in S}A_i + \frac{f_{m+1}}{(\sum A)-m}\]
後は \(\sum_{|S|=m}\prod_{i\in S}A_i\) を高速に計算すればよいです。これは \(\prod (1+A_i x)\) の \(x^m\) の係数なので、分割統治法により \(m=0,\dots,N\) の値を \(O(N \log^2 N)\) で計算できます。
以上より、全体で \(O(N \log^2 N)\) の計算量で解くことができました。
posted:
last update: