Official

Ex - Trio Editorial by Nyaan


この問題は 多項式の逆元 の練習問題です。多項式・形式的冪級数の逆元の計算に帰着できる数え上げは各所で出題されているので、数え上げの知識をつけたい人はこの問題で練習してみましょう。

まず、\(f(t)\) を次のように定めます。求めたい答えは \(f(T)\) です。

  • \(f(t)\) \(:=\) (時刻 \(0\) に 3 人がそれぞれ \(A, B, C\) にいて、時刻 \(t\)はじめて 同じ地点に集まる確率)

条件の「はじめて」という部分が厄介なので、まずはこれを取り除いた場合を数えて、そこから \(f(t)\) を計算する、という手順で計算してみましょう。
\(g(t), h(t)\) を次のように定めます。

  • \(g(t)\) \(:=\) (時刻 \(0\) に 3 人がそれぞれ \(A, B, C\) にいて、時刻 \(t\) に同じ地点に集まる確率)

  • \(h(t)\) \(:=\) (時刻 \(0\) に 3 人がそれぞれ \(0, 0, 0\) にいて、時刻 \(t\) に同じ地点に集まる確率)

あらかじめ \(g(t), h(t)\) \((0 \leq t \leq T)\) が計算出来ているとします。すると、\(f(T)\) は次の \(f(t), g(t), h(t)\) 間の関係式を用いて \(\mathrm{O}(T^2)\) で計算できます。(俗に「除原理」と呼ばれている DP です)

\[f(t) = g(t) - \sum_{0 \leq u \lt t} f(u) h(t - u)\]

この DP を高速化しましょう。まず、上式は右辺が \(f\) に関する畳み込みの形をしているので、ABC212 の H 問題の解説 でも説明した分割統治 FFT によって \(\mathrm{O}(T \log^2 T)\) で計算できます。この解法でも AC するのに十分高速ですが、母関数を利用するとさらに高速化できます。
\(f, g, h\) の通常型母関数 \(F, G, H\) は次の通りです。

\[F(x) = \sum_{i=0}^\infty f(i) x^i, G(x) = \sum_{i=0}^\infty g(i) x^i, H(x) = \sum_{i=0}^\infty h(i) x^i\]

\(F\)\(G, H\) の関係式で表してみましょう。これは先の漸化式を利用すれば、

\[ \begin{aligned} &\sum_{t=0}^\infty f(t) x^t = \sum_{t=0}^\infty \left( g(t) - \sum_{0 \leq u \lt t} f(u) h(t - u) \right) x^t \\ &\to F(x) = G(x) - F(x) (H(x) - h(0)) \end{aligned} \]

で、\(h(0) = 1\) を利用すると、

\[ \begin{aligned} &\to F(x) = G(x) - F(x) (H(x) - 1) \\ &\to F(x) H(x) = G(x) \\ &\to F(x) = \frac{G(x)}{H(x)} \end{aligned} \]

と表せます。よって \(H(x)\) の乗法的逆元 \(\frac{1}{H(x)}\)\(T\) 次の係数まで計算できれば \(F(x)\)\(T\) 次の係数を計算できます。形式的冪級数の逆元は ABC260 Ex の解説 でも書いた通り \(\mathrm{O}(T \log T)\) で計算できます。

以上の考察により、\(g, h\) が計算できれば \(f(T)\)\(\mathrm{O}(T \log T)\) で計算できるとわかりました。\(g, h\) を計算するために次の部分問題を解いていきましょう。

時刻 \(0\) に 3 人がそれぞれ \(X, Y, Z\) にいます。時刻 \(t\) に3 人が同じ地点にいる確率を \(p(t)\) とします。 \(p(0), p(1), \dots, p(T)\) を列挙してください。

(以下では \(K\) が負である場合や整数でない場合に \(\frac{1}{K!} = 0, \binom{N}{K} = 0\) と定義します。)

\(t\) 秒後に地点 \(w\) に 3 人が集まる確率を \(p(t, w)\) とします。\(p(t, w)\)

\[p(t, w) = \frac{1}{2^{3t}} \binom{t}{\frac{t + w - X}{2}} \binom{t}{\frac{t + w - Y}{2}} \binom{t}{\frac{t + w - Z}{2}}\]

と表せます。

\[q(i) = \frac{1}{\left(\frac{i-X}{2}\right)! \left(\frac{i-Y}{2}\right)!\left(\frac{i-Z}{2}\right)!}\]

\[r(i) = \frac{1}{\left(\frac{i+X}{2}\right)! \left(\frac{i+Y}{2}\right)!\left(\frac{i+Z}{2}\right)!}\]

とすると

\[p(t, w) = \frac{(t!)^3}{2^{3t}} q(t + w) r(t - w)\]

と表せるので、\(t + w = u\) と置換して

\[p(t) = \sum_w p(t, w) = \frac{(t!)^3}{2^{3t}} \sum_u q(u) r(2t - u)\]

という式を得られます。よって、二項係数を前計算して \(q, r\) の列を生成して、\(q\)\(r\) を畳み込んた列の \(2t\) \((0 \leq t \leq T)\) 番目の値を取り出して \(\frac{(t!)^3}{2^{3t}}\) 倍すれば \(p(t)\)\(0 \leq t \leq T\) について列挙できます。(\(u\)\(2t - u\) が負の index を取り得るので、畳み込みの際に適宜 offset を取るなどの工夫が必要なのに注意してください。)

以上よりこの問題を \(\mathrm{O}((A+B+C+T) \log (A+B+C+T))\)\(\mathrm{O}(T \log T + A + B + C)\) で計算出来て、 これは十分高速です。(逆元の計算に \(\mathrm{O}(T \log^2 T)\) かける解法でも AC できると思います。)

Bonus : 人数を \(3\) 人から \(N\)\((N \leq 10^5)\) に一般化した場合を解いてみましょう。

posted:
last update: