G - Colorful Candies 2 Editorial by noshi91


公式解説の内容と記法を前提として、時間計算量 \(\mathrm{O} (N \log (N) + \log(P))\) の解法を説明します。 (長さ \(N\)\(\mathbb{Z} / P \mathbb{Z}\) の元の列の畳み込みが時間計算量 \(\mathrm{O} (N \log (N) + \log(P))\) で行えると仮定しています。)

\[\begin{aligned} \mathbb{E} \left \lbrack \sum _ {i = 1} ^ {C} X_i \right \rbrack & = \sum _ {i = 1} ^ {C} \mathbb{E} \left \lbrack X_i \right \rbrack \cr & = \sum _ {i = 1} ^ {C} \frac {\binom{N}{K} - \binom{N - n _ i}{K}} {\binom{N}{K}} \cr & = C - \binom{N}{K} ^ {-1} \sum _ {i = 1} ^ {C} \binom{N - n _ i}{K} \end{aligned}\]

したがって、各 \(K\) について \(\displaystyle \sum _ {i = 1} ^ {C} \binom{N - n _ i}{K}\) を計算できればよいです。

ここで \(f _ j = \# \left \lbrace i ~ \vert ~ N - n _ i = j \right \rbrace\) と定義すれば

\[\begin{aligned} \sum _ {i = 1} ^ {C} \binom{N - n _ i}{K} & = \sum _ {j = K} ^ {N} f _ j \binom{j}{K} \cr & = \sum _ {j = K} ^ {N} \frac{f _ j ~ j!}{K!(j-K)!} \cr & = \frac{1}{K!} \sum _ {\substack{K \leq j \leq N \cr j + l = N + K}} \frac{f _ j ~ j!}{(N - l)!} \end{aligned}\]

最後の式は概ね畳み込みの式になっています。 まず、\(0 \leq l \leq N\) が常に満たされるのでそれを制約に加えても問題ありません。 更に \(l \leq N, ~ j + l = N + K\) から \(K \leq j\) は導かれるので、\(K \leq j\)\(0 \leq j\) に緩めても問題ありません。

\[\begin{aligned} s _ j & = f _ j ~ j ! \cr t _ l & = (N - l)! ^ {-1} \end{aligned}\]

と定義することで

\[\begin{aligned} \sum _ {\substack{K \leq j \leq N \cr j + l = N + K}} \frac{f _ j ~ j!}{(N - l)!} & = \sum _ {\substack{0 \leq j \leq N \cr 0 \leq l \leq N \cr j + l = N + K}} s _ j ~ t _ l \end{aligned}\]

を得ます。 \(s, t\) を畳み込むことで全ての \(K\) について上の式の値が求まります。 時間計算量は \(\mathrm{O} (N \log (N) + \log(P))\) です。

実装例 (簡略化のため、時間計算量が \(\Theta(N (\log(N) + \log(P)))\) に悪化しています。)


以上のアルゴリズムは、\(\displaystyle \sum _ {i = 1} ^ {C} (1 + x) ^ {N - n _ i}\) の係数を計算するために \(\displaystyle \sum _ {i = 1} ^ {C} x ^ {N - n _ i}\)Taylor Shift していると解釈することもできます。

posted:
last update: