この問題は power projection と呼ばれるアルゴリズムをテーマにした問題です。power projection は FPS 24 題 では扱いませんでしたが、汎用性が高いことから 2025 年の高難度数え上げ (特に AGC や企業コン、Universal Cup) で適用可能な問題が多く出題されたテクニックです。高難度問題に挑戦したい方は是非マスターしてください。

power projection

この章は ABC387-G 解説 に加筆したものです。

power projection と呼ばれる問題について説明します。具体的には次の問題が power projection です。

power projection

\(n\) 次の多項式 \(f(x), g(x)\) が与えられる。

\[h_i = \lbrack x^n \rbrack f(x)^i g(x)\]

としたときに \(h_0, h_1, \dots, h_n\) を列挙せよ。

power projection は 2024 年に noshi91 さんによって \(\mathrm{O}(n \log^2 n)\) で解く方法が発見されました。(初出記事)
まず、\(\lbrack x^0 \rbrack f(x) \neq 0\) の場合を考えます。\(\lbrack x^0 \rbrack f(x) = c\) とおくと

\[\lbrack x^k \rbrack f(x)^i g(x) = \sum_{j = 0}^i \binom{i}{j} c^j \lbrack x^k \rbrack (f(x)-c)^{i-j} g(x)\]

となるので \(f(x) \gets f(x)-c\) の場合の答えから畳み込みで求まります。よって \(\lbrack x^0 \rbrack f(x) = 0\) の場合が解ければよいです。以降では \(\lbrack x^0 \rbrack f(x) = 0\) とします。

求めたい答えを 2 変数形式的冪級数で表すと、以下の値が \(y\) の多項式として \(n\) 次まで求まればよいです。

\[[x^n] \sum_{i=0}^\infty y^i f(x)^i g(x) = \lbrack x^n \rbrack \frac{g(x)}{1 - y f(x)}\]

この式に Bostan-Mori 法を適用します。(Bostan-Mori 法については ABC300-Ex 解説 を参照してください。)すると操作ごとに \(y\) の次数は倍々になっていきますが、\(x\) の次数は半分ずつになっていきます。すなわち \(t\) 回目のイテレーションでは \(x\) は \(\lfloor n/2^t \rfloor\) 次、 \(y\) は \(2^t\) 次です。よって各イテレーションでの多項式乗算の計算量は \(\mathrm{O}(n \log n)\) になり、全体で \(\mathrm{O}(\log n)\) 回の乗算が必要なので計算量は \(\mathrm{O}(n \log^2 n)\) となります。

実際にアルゴリズムを実装したコードが以下になります。

• Nyaan’s Library の形式的冪級数ライブラリを使用したコードとなっています。pre(n) は fps 型の要素の先頭 n 項を取り出す関数です。それ以外に分からない関数定義があれば適宜ライブラリを参照ください。
• 説明用のため定数倍高速化を行っておらず低速である点に注意してください。

// [x^n] f(x)^i g(x) を i=0,1,...,n で列挙
// 簡単のため [x^0] f = 0 を仮定
template <typename mint>
FormalPowerSeries<mint> power_projection(FormalPowerSeries<mint> f,
                                         FormalPowerSeries<mint> g) {
  using fps = FormalPowerSeries<mint>;
  assert(f.size() == g.size());
  int n = f.size() - 1, k = 1;
  vector<fps> P(n + 1, fps(k + 1));
  vector<fps> Q(n + 1, fps(k + 1));
  Q[0][0] = 1;
  for (int i = 0; i <= n; i++) P[i][0] = g[i];
  for (int i = 0; i <= n; i++) Q[i][1] = -f[i];
  while (n) {
    auto R = Q;
    for (int i = 1; i <= n; i += 2) R[i] = -R[i];
    auto S = multiply2d(P, R);
    auto T = multiply2d(Q, R);
    vector<fps> U(n / 2 + 1, fps(k * 2 + 1));
    vector<fps> V(n / 2 + 1, fps(k * 2 + 1));
    for (int i = 0; i <= n / 2; i++) U[i] = S[i * 2 + n % 2], V[i] = T[i * 2];
    P = U, Q = V;
    n /= 2, k *= 2;
  }
  return (P[0] * Q[0].inv()).pre(f.size());
}

解法

はじめに \(f_n, g_n\) を次のように置きます。

• \(f_n :=\) (ある 1 人が \(n\) 回行動した時にまだゴールしてしない確率)
• \(g_n :=\) (ある 1 人が \(n\) 回行動した時に \(n\) 回目ではじめてゴールする確率)

明らかに \(g_n = f_{n-1} - f_n\) \((n \geq 1)\) かつ \(n \geq N\) の時 \(f_n = 0\) です。よって \(f_0, f_1, \dots, f_{N-1}\) を計算できればよいです。

\(D(x)\) を「サイコロを \(1\) 回振った時に \(n\) マス進む確率」の母関数とします。すなわち

\[D(x) = \frac{1}{M} \sum_{i=1}^M x^{A_i}\]

です。このとき、サイコロを \(n\) 回振ってマス \(k\) \((0 \leq k \lt N)\) にいる確率は \([x^k] D^n\) と表せます。よって

\[ \begin{aligned} f_n &= \sum_{k=0}^{N-1} [x^k] D^n \\ &=[x^{N-1}] (1+x+\dots + x^{N-1}) D^n \end{aligned} \]

と表すことが出来ます。上式を見ると \([x^{\text{const.}}] F^n G\) という式の形をしているため power projection がそのまま適用できます。つまり、\(f_0, f_1, \dots, f_{n-1}\) の列挙が \(\mathrm{O}(N \log^2 N)\) でできます。

\(f_n, g_n\) が全て計算できたのでこれをもとに答えを求めましょう。人 \(i\) が \(k\) 回目の行動でゴールして勝つ確率は

\[g_k f_{k}^{i-1} f_{k-1}^{L-i}\]

です。これを \(k=1,2,\dots,N\) について足し合わせたものが人 \(i\) に対する答えとなります。
\(i=L\) の場合は上式を愚直に足し合わせて計算することにして、\(1 \leq i \lt L\) とします。すると \(f_{k-1} = 0\) の時に上式は \(0\) になるので無視できて、\(f_{k-1} \neq 0\) を仮定できます。すると上式は

\[ \begin{aligned} &g_k f_{k-1}^{L-1} \left(\frac{f_{k}}{f_{k-1}}\right)^{i-1} \\ &= [x^{i-1}] \frac{g_k f_{k-1}^{L-1} }{1 - \frac{f_{k}}{f_{k-1}} x} \end{aligned} \]

と変形できます。これを \(k=1,2,\dots,N\) について足し合わせた

\[\sum_{1 \leq k \leq N, f_{k-1} \neq 0} [x^{i-1}] \frac{g_k f_{k-1}^{L-1} }{1 - \frac{f_{k}}{f_{k-1}} x} = [x^{i-1}] \sum_{1 \leq k \leq N, f_{k-1} \neq 0} \frac{g_k f_{k-1}^{L-1} }{1 - \frac{f_{k}}{f_{k-1}} x}\]

が人 \(i\) に対する答えです。よって

\[\sum_{1 \leq k \leq N, f_{k-1} \neq 0} \frac{g_k f_{k-1}^{L-1} }{1 - \frac{f_{k}}{f_{k-1}} x}\]

の \(0\) 次から \(L-2\) 次までの係数を求めればよく、これは FPS のマージテクと俗に呼ばれる、セグ木状に分数を足し合わせていく手法で分数の分子と分母を \(\mathrm{O}(N \log^2 N)\) で計算した後に多項式の逆元 (多項式の逆元は ABC260-Ex 解説 を参照してください) で \(\mathrm{O}(L \log L)\) 掛けて分母の逆元を求めて分子とかけ合わせれば係数を求められます。

以上よりこの問題を \(\mathrm{O}(N \log^2 N + L \log L)\) 程度で解くことが出来て、これは十分高速です。