G - Colorful Candies 2 Editorial by tsutaj

高橋君がもらう \(K\) 個のキャンディに含まれる色の種類数の期待値 (= 答え) に \(\binom{N}{K} = {}_N C_K\) を掛けた値は、高橋君がもらうキャンディに含まれる色の種類数を全ての選び方について求め、その総和をとって得られる値と等しいです。後者の「総和」を求める \(O(N \log^2 N)\) のアルゴリズムを解説します。

キャンディの色が 1 種類のみの場合

キャンディが \(N\) 個あり、それらの色は全て同じであるとします。ここで、\(0 \leq k \leq N\) について 2 種類の変数を以下のように定義します。

  • \(\mathrm{dp}_k := N\) 個のキャンディから \(k\) 個もらう場合の数
  • \(\mathrm{sum}_k := \) すべての「\(N\) 個のキャンディから \(k\) 個もらう」場合について、もらうキャンディに含まれている色の種類数を求め、その総和を取って得られる値

仮定より色が 1 種類しかないため、それぞれの変数の値は次のように求められます。

  • \(\mathrm{dp}_k = \binom{N}{k}\)
  • \(\mathrm{sum}_k = \begin{cases} 0 & k = 0 \\ \mathrm{dp}_k & \mathrm{otherwise} \end{cases}\)
    • \(k = 0\) ならばキャンディをひとつももらっていないため、種類数は \(0\)
    • それ以外の場合、種類数は \(1\) なので、総和は \(\mathrm{dp}_k\) と等しい

キャンディの色が複数種類ある場合

色 \(1\) から色 \(C\) までのキャンディのみが存在すると仮定します。 キャンディの並び順は考慮する必要がなく、各色のキャンディが何個ずつあるかのみによって答えが定まります。\(i = 1, 2, \ldots, C\) について、色 \(i\) のキャンディの個数を \(n_i\) で表します。

それぞれの色について、「キャンディの色が 1 種類のみの場合」で前述した \(\mathrm{dp}, \mathrm{sum}\) の値を求めておきます。つまり各色 \(i\) について、キャンディの色の集合が \(\left\{ i \right\}\) であるときの値を求めます。これらの情報を併合することで、答えを求めることを考えます。

キャンディの色の集合 \(S, T\) \((S \cap T = \emptyset)\) それぞれにおける \(\mathrm{dp}, \mathrm{sum}\) の値を以下のように併合することによって、キャンディの色の集合 \(S \cup T\) における \(\mathrm{dp}, \mathrm{sum}\) の値を求めることができます。なお \(\mathrm{dp}_{X, k}\) とは、キャンディの色の集合が \(X\) であるときの \(\mathrm{dp}_k\) の値であり、\(\mathrm{sum}_{X, k}\) も同様です。また、\(N_X = \sum_{i \in X} n_i\) です。

  • \(\mathrm{dp}_{S \cup T, k} = \sum_{i+j = k} \mathrm{dp}_{S, i} \times \mathrm{dp}_{T, j}\)
  • \(\mathrm{sum}_{S \cup T, k} = \sum_{i+j = k} \left\{ \mathrm{sum}_{S, i} \times \mathrm{dp}_{T, j} + \mathrm{dp}_{S, i} \times \mathrm{sum}_{T, j} \right\}\)
    • \(k\) は \(0 \leq k \leq N_S + N_T\) の範囲を取りうる

この計算は、高速フーリエ変換で \(\left( \mathrm{dp}_{S}, \mathrm{dp}_{T} \right)\) と \(\left( \mathrm{sum}_S, \mathrm{dp}_T \right)\) と\(\left( \mathrm{dp}_S, \mathrm{sum}_T \right)\) をそれぞれ畳み込むことによって高速に行えます。

これを繰り返すことによって、集合 \(\left\{ 1, 2, \ldots, C \right\}\) における \(\mathrm{dp}, \mathrm{sum}\) の値を求めることができ、\(\mathrm{sum}\) を利用して元の問題の答えも求められます。集合を併合するときは、ある集合 \(X\) に属しているキャンディの個数 \(N_X\) が少ない順に集合を 2 つ選択してそれを併合することにすると、全体の処理が高速に行えます。計算量は \(O(N \log^2 N)\) です。

実装例

posted:
last update: