Editorial - AtCoder Beginner Contest 387

Contest Duration: 2025-01-04 21:00:00+0900 - 2025-01-04 22:40:00+0900 (local time) (100 minutes) Back to Home

Official

G - Prime Circuit Editorial by Nyaan

今回は競技プログラミング界において 2024 年最大の発見である Kinoshita-Li 合成 をテーマとした出題です。

形式的冪級数の合成

形式的冪級数は競技プログラミングにおいて非常に強力なツールです。これは過去の ABC,ARC,AGC の出題例を見てもわかると思います。また、競技プログラミングに限らず、組合せ数学や暗号理論の分野でも形式的冪級数は理論の礎になっています。今回はそういった形式的冪級数の実用を支える基本的な演算に注目してみましょう。
2 つの形式的冪級数を “マージ” する方法は大きく分けて 3 種類あり、和 \(f(x) + g(x)\) ・積 \(f(x) g(x)\) ・合成 \(g(f(x))\) です。それぞれの計算量を考えると、一般的な知識では

和 \(f(x) + g(x)\) : (次数を \(n\) として) \(\mathrm{O}(n)\)
積 \(f(x) g(x)\) : FFT を利用して \(\mathrm{O}(n \log n)\)
合成 \(g(f(x))\) : ???

という感じになるのではないかと思います。今回は合成 \(g(f(x))\) に注目します。
関数合成の教科書的なアルゴリズムは、 Brent と Kung によって発見された \(\mathrm{O}((n \log n)^{3/2})\) のアルゴリズムでした。(論文) このアルゴリズムは計算量が示す通り和や積に比べてはるかに低速であるため、競技プログラミングの世界において関数合成が用いられることはほとんどありませんでした。
しかし、2024 年、Kinoshita(noshi91) と Li(Elegia) はこの関数合成を \(\mathrm{O}(n \log^2 n)\) で行うアルゴリズムを発見しました。(論文) このアルゴリズムは画期的で、実用上も高速に動作するため、瞬く間に競技プログラミング界隈(の一部)に広がりました。(以降ではこのアルゴリズムを Kinoshita-Li 合成と呼びます。)
今回はこの Kinoshita-Li 合成の仕組みに関する解説を行います。

予備知識：転置原理

Kinoshita-Li 合成を知る上で必要な知識は 2 つあります。1 つは Bostan-Mori 法、もう 1 つは 転置原理 です。Bostan-Mori 法は ABC300-Ex 解説で詳しく触れたのでこの解説では説明しません。転置原理について解説します。

転置原理については Nachia さんの記事が分かりやすいのでそちらも合わせてお読みください。

概要

転置原理の転置とは、行列の転置のことです。\(N \times M\) 行列 \(A\) の転置とは、\((i, j)\) 成分を \((j, i)\) に移動して出来る \(M \times N\) 行列のことを言い、\(A^T\) と表します。例えば

\[ \left( \begin{array}{ccc} 1 & 2 & 3 \\ 4 & 5 & 6 \end{array} \right)^T = \left( \begin{array}{cc} 1 & 4 \\ 2 & 5 \\ 3 & 6 \end{array} \right) \]

です。(逆行列と混同しないようご注意ください)

転置原理とはおおよそ次のような原理です。

転置原理

\(n \times m\) 行列 \(A\) がある。
今、長さ \(m\) のベクトル \(v\) が与えられたときに長さ \(n\) のベクトル \(w = Av\) を計算量 \(X\) で計算するアルゴリズムがあるとする。
あなたは長さ \(n\) のベクトル \(w\) が与えられたときに長さ \(m\) のベクトル \(v = A^T w\) を計算するアルゴリズムを知りたい。
実は \(A^T w\) は計算量 \(X\) で計算することが出来て、その手順も機械的に求めることが出来る。(!!)

アルゴリズムの構成の仕方

転置原理を用いた \(A^T w\) を計算するアルゴリズムの構成の仕方は次の通りです。
\(A\) を計算するアルゴリズムを 1 個 1 個の手順に分解したものが、行列積を用いて \(A = A_1 A_2 \dots A_N\) と表せるとします。この時 \(A^T = A_N^T \dots A_2^T A_1^T\) です。よって、

\(A\) を行列積に分解する : \(A_1, A_2, \dots, A_N\)
逆向きに並べ替える : \(A_N, \dots, A_2, A_1\)
作用を転置させる : \(A_N^T, \dots, A_2, A_1\)

というように変形した上で、変形後の行列積に対応するアルゴリズムをこの順で作用させることで \(A^T w\) を計算できます。

操作の分解と転置

上で \(A\) を \(A_1, A_2, \dots, \) と分解しましたが、この分解の最小単位がどのような構造をしているのかを考えます。
最小単位は基本的には次の 3 つです。\(x, y\) を変数(ベクトルの成分)、\(m\) を定数としたときに

\(x, y\) を swap
\(x \gets m \times x\)
\(x \gets x + m \times y\)

線形代数の知識がある方はご存じの通り、アルゴリズムの計算過程で登場する変数があまねく入力されたベクトルの線形結合であることは、アルゴリズムが上記の 3 つの操作のみを用いて記述できることと同値です。よって、アルゴリズムが線形変換であればこれらの操作で書き表せることになります。
3 つの操作を転置すると次のようになります。このようになることは操作に対応する行列の転置を考えると分かりやすいと思います。

\(x, y\) を swap : 転置後そのまま
\(x \gets m \times x\) : 転置後そのまま
\(x \gets x + m \times y\) : 転置後 \(y \gets y + m \times x\)

ただし、転置原理を実際にアルゴリズムに適用する際は、考察を楽にするためにもう少し大きいカタマリ(行列)に分解して転置を考えることが多いです。以下、 maspy さんの記事より引用します。

\(F\) のアルゴリズムから \(F^{T}\) を導出することは， \(F\) のアルゴリズムの算術演算を全て列挙していけば原理的にはできるはずですが，アルゴリズムが複雑だとそのような方法での機械的なアルゴリズム導出は大変です．実際には，算術演算単位ではなくもう少し大きなステップのアルゴリズムに分解して，各ステップの転置問題を（転置原理を用いずに）解くことで導出することが多いと思います．

例：多項式乗算の転置

\(A(x)\) を \(n\) 次の多項式, \(B(x)\) を \(m\) 次の多項式とします。\(C(x) = A(x) \times B(x)\) を \(\mathrm{O}(nm)\) で計算するアルゴリズムの転置を考えます。
ここで、\(A, B\) を両方とも変数(=入力で与えられるベクトル)とみなすと、計算過程が行列の作用という形で表せず転置できません。そこで \(A\) を変数(ベクトルに対応)、\(B\) を定数(行列に対応)として考えることにします。すなわち

\[a = \left( \begin{array}{c} A_0 \\ A_1 \\ A_2 \\ \vdots \\ A_n \end{array} \right), c = \left( \begin{array}{c} C_0 \\ C_1 \\ C_2 \\ \vdots \\ C_{n+m} \end{array} \right) \]

\[ B = \left( \begin{array}{cccc} B_0 & 0 & \dots & 0 & 0 \\ B_1 & B_0 & \dots& 0 & 0 \\ \vdots & & \ddots & & \vdots \\ 0 & 0 & \dots& B_m & B_{m-1} \\ 0 & 0 & \dots& 0 & B_m \\ \end{array} \right) \]

としたときに \(c \gets Ba\) を計算するアルゴリズムを考えます。
転置前のアルゴリズムは次の通りです。

for(int k = 0; k <= n + m; k++) {
  for(int i = max(0, k - m); i <= min(k, n); i++) {
    c[k] += a[i] * b[k - i];
  }
}

これを先の手順に従って機械的に転置すると次のようになります。

for(int k = n + m; k >= 0; k--) {
  for(int i = min(k, n); i >= max(0, k - m); i--) {
    a[i] += c[k] * b[k - i];  
  }  
}

このアルゴリズムは確かに \(a \gets B^T c\) を計算するアルゴリズムになっていることが確認できると思います。

また、 \(c \gets B^T a\) を計算する転置前のアルゴリズムは FFT を用いて \(\mathrm{O}((n+m) \log (n+m))\) になることが知られています。この FFT を利用した畳み込みアルゴリズムを最小単位に分解して転置することで、\(a \gets B^T c\) を計算する転置後のアルゴリズムも \(\mathrm{O}((n+m) \log (n+m))\) で計算することが出来ることがわかります。(実際に FFT を最小単位に分解するのは骨が折れるので、転置を行う際は FFT を 1 つのカタマリとして転置を行うことになります。)

転置原理に限らない一般的な話として、\(B^T c\) の計算は \(b\) を reverse して FFT すれば \(\mathrm{O}((n+m) \log (n+m))\) で計算できるというのは有名な話かと思います。ところが実は、転置原理を経由すると最大で 2 倍定数倍の良いアルゴリズムを得られることが出来て、 middle product と呼ばれています。

ここまでの話だと内容の捉え所が難しい面があるように思っています。(主に「転置原理を何に使うの？」といった意味で) ですが、転置原理を理解することで非自明なアルゴリズムを導出することが出来て、それが Kinoshita-Li 合成の発見につながりました。

転置原理の説明は以上です。予備知識の説明は終えたので、本題である Kinoshita-Li 合成の説明をしていこうと思います。

power projection (Kinoshita-Li 合成の転置)

はじめに power projection と呼ばれる問題について説明します。(このアルゴリズムは実は Kinoshita-Li 合成の転置になっています。)

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 は \(\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 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 法を適用します。すると操作ごとに \(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());
}

Kinoshita-Li 合成

power projection を転置してみることにします。power projection は \(a_{i,j} = [x^j] f^i, v = \mathrm{rev}(g)\) と置いたときに \((a_{i,j}) v\) を計算するアルゴリズムです。では、転置後のアルゴリズムで得られるベクトルの第 \(k\) 成分が何になるかを考えてみます。

\[ \begin{aligned} \left((a_{i, j})^T w\right)_k &= \left((a_{j,i})w\right)_k \\ &= \sum_{i=0}^n [x^k] f^i w_i \\ &= [x^k] \sum_{i=0}^n w_i f^i \end{aligned} \]

\(g(x) = \sum_{i=0}^n w_i x^i\) とおくと

\[= [x^k] g(f(x))\]

となります。よって power projection の転置は関数合成になることがわかりました！
よって、power projection のアルゴリズムを転置することで合成の \(\mathrm{O}(n \log^2 n)\) アルゴリズムを得られることになります。転置前のアルゴリズムは「代入などの基本操作」+ 「畳み込み」+ 「2 次元畳み込み」で構成されています。畳み込みの転置は先に触れたように reverse + 畳み込みになり、2 次元畳み込みも同様に reverse + 畳み込みになります。よって転置を利用して機械的に合成のアルゴリズムを記述できます。

以上の内容を適切に実装したものが以下のコードです。このコードが power projection の転置になっていることを理解することは容易ではないかと思いますが、変数の流れを丁寧に観察すると転置になっていることが確認できると思います。

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

template <typename mint>
vector<FormalPowerSeries<mint>> inner(const FormalPowerSeries<mint>& g,
                                      vector<FormalPowerSeries<mint>> Q, int n,
                                      int k) {
  using fps = FormalPowerSeries<mint>;
  if (n == 0) {
    fps h = g * Q[0].inv().rev();
    vector<fps> P(n + 1, fps(k + 1));
    for (int i = 0; i < (int)g.size(); i++) P[0][i] = h[i + Q[0].size() - 1];
    return P;
  }
  auto R = Q;
  for (int i = 1; i <= n; i += 2) R[i] = -R[i];
  auto T = multiply2d(Q, R);
  vector<fps> V(n / 2 + 1, fps(k * 2 + 1));
  for (int i = 0; i <= n / 2; i++) V[i] = T[i * 2];
  auto U = inner(g, V, n / 2, k * 2);
  vector<fps> S(n * 2 + 1, fps(k * 2 + 1));
  for (int i = 0; i <= n / 2; i++) S[i * 2 + n % 2] = U[i];
  vector<fps> revR(n + 1, fps(k + 1));
  for (int i = 0; i <= n; i++) revR[i] = R[n - i].rev();
  vector<fps> Pshift = multiply2d(S, revR);
  vector<fps> P(n + 1, fps(k + 1));
  for (int i = 0; i <= n; i++) {
    for (int j = 0; j <= k; j++) P[i][j] = Pshift[i + n][j + k];
  }
  return P;
}

// g(f(x)) mod x^{n+1} を計算
// 簡単のため [x^0] f(x) = 0 を仮定
template <typename mint>
FormalPowerSeries<mint> composite(FormalPowerSeries<mint> f,
                                  FormalPowerSeries<mint> g) {
  using fps = FormalPowerSeries<mint>;
  assert(f.size() == g.size());
  int n = f.size() - 1, k = 1;
  vector<fps> Q(n + 1, fps(k + 1));
  Q[0][0] = 1;
  for (int i = 0; i < n + 1; i++) Q[i][1] = -f[i];
  auto P = inner(g, Q, n, k);
  fps h(n + 1);
  for (int i = 0; i <= n; i++) h[i] = P[i][0];
  return h.rev();
}

Kinoshita-Li 逆関数

Kinoshita-Li 合成の副産物として、\(f(x) \bmod x^{n+1}\) \(([x^0]f(x) = 0, [x^1]f(x) \neq 0)\) が与えられたとき逆関数 \(f^{\langle -1 \rangle}(x) \bmod{x^{n+1}}\) を計算するアルゴリズムがあります。今回の問題はこれを利用して解くこともできるのでこちらも解説します。

ラグランジュの反転公式より(詳しくは ABC345G 解説を参照ください)

\[[x^n] f(x)^k = \frac{k}{n} [x^{n-k}] \left( \frac{x}{f^{\langle -1 \rangle}(x)}\right)^n\]

が成り立ちます。よって \(f(x)\) に対して power projection を適用して得られた列に適宜定数倍を掛けて reverse すれば \(\left( \frac{x}{f^{\langle -1 \rangle}(x)}\right)^n \bmod x^{n+1}\) が得られます。あとは形式的冪級数の \(\log, \exp\) を適切に利用すれば \(f^{\langle -1 \rangle}(x)\) を計算できます。(詳しくは Nyaan による実装例を参照ください)

問題の解法

さて、問題の解法に移ります。
まず、問題文の条件を数えやすい形に言い換えます。グラフ \(G\) が問題文の条件を満たすことは以下の条件が成り立つことと同値です。

\(G\) に含まれる全ての閉路同士が頂点を共有せず、かつ全ての閉路の長さが素数である。\(\cdots (\ast)\)

(証明) \((\ast)\) が成り立てば問題文の条件が成り立つことは明らかであるから、問題文の条件が成り立つ時に \((\ast)\) が成り立つことを言えばよい。ここで次の補題を挙げる。

補題 : グラフ \(G\) が閉路 \(c_1, c_2\) を含み、かつ \(c_1\) と \(c_2\) が少なくとも頂点を共有するとする。この時、\(G\) は偶数長の回路を含む。

補題が示せれば対偶を取って適切に議論することで示したい命題を示せる。よって補題が証明できればよい。以降では補題を証明する。
\(c_1, c_2\) が偶数長の場合は明らかに成り立つので \(c_1, c_2\) の長さが奇数である場合のみを考える。

(i) \(c_1, c_2\) が辺を共有しない場合

この場合は \(c_1, c_2\) の全ての辺を \(1\) 回ずつ通る回路が存在して、この回路の長さは偶数であるから補題は成り立つ。

(ii) \(c_1, c_2\) が辺を共有する場合

\(c_1\) と \(c_2\) の丁度一方にのみ含まれる辺からなるグラフを考える。このグラフの全ての頂点の次数は偶数であるから、いくつかの閉路からなることがわかる。閉路の個数が \(1\) 個の場合は、閉路の長さが偶数になるため補題が成り立つ。閉路の個数が \(2\) 個以上の場合を考える。閉路のうち長さが最小のものを取り出し \(c_3\) とする。
\(c_3\) の長さが偶数の場合は補題が成り立つ。\(c_3\) の長さが奇数の場合を考える。\(|c_1| \geq |c_2|\) として一般性を失わない。\(c_1\) と \(c_2\) の丁度一方にのみ含まれる辺からなるグラフの辺数は \(|c_1| + |c_2| - 2\) 以下であるから \(|c_3| \leq (|c_1| + |c_2| - 2) / 2\lt |c_1|\) である。また、\(c_2\) と \(c_3\) は明らかに辺を共有するので、 \(c_2\) と \(c_3\) に注目することでよりサイズの小さい問題に帰着される。

以上の議論より、\(c_1, c_2\) の辺数の合計に対して帰納法を回すことで補題を証明できることが確認できる。よって命題は示された。(証明終わり)

さて、以上の証明により結局次の問題が解ければよいです。

言い換え後の問題

\(N\) 頂点の単純連結グラフ \(G\) のうち次の条件を満たすものの個数は？

\(G\) に含まれる全ての閉路同士が頂点を共有せず、かつ全ての閉路の長さが素数である。

この問題を母関数を用いて解くことを考えます。
この手の問題を解くときの典型的な方針として、「条件を満たすグラフを 1 つ取ってきて、さらに 1 頂点を根とみなした通り数」の数列の指数型母関数を考えます。すなわち、問題の答えを \(a_N\) としたときに

\[F(x) = \sum_{n=0}^\infty \frac{n a_n}{n!} x^n\]

を満たす母関数 \(F(x)\) を考えます。

同種の言い換えに「頂点に \(0, 1, \dots, n-1\) のラベルを張って頂点 \(0\) を根とする根付き木とみなす」という言い換え方があります。この言い換えでは \(n\) 頂点の木に \(1\) から \(n-1\) までのラベルを自由に貼れることから、母関数を考える時は \(\sum_n \frac{a_n}{(n-1)!} x^n\) というような母関数を考えるとキレイに計算できることが多いです。この式は結果的に上式と一致して、以降同様の議論を行えます。

さて、この時「根にいくつかの子が付いた状態」を表す母関数は \(x \exp(F(x))\) になります。また、根を含む頂点がサイズ \(p\) の閉路をなしている場合は根を基準に \(x \exp(F(x))\) が列として \(p\) 個並んでいることになるので \((x \exp(F(x)))^p\) になります。ただし 1 つの閉路に 2 つの列が対応しているので \(2\) で割る必要があります。よって、

\[G(x) = x + \sum_{p \geq 3, p \text{ is prime}} \frac{x^p}{2}\]

と置いた時 \(F(x)\) は

\[F(x) = G(x \exp F(x))\]

と表すことが出来ます。
この式を利用してニュートン法を行うことでこの問題を解くことが出来ます。(ニュートン法については ABC345G 解説を参考にしてください) すなわち、

\[A(F) = G(x \exp(F(x))) - F(x) = 0\]

と置くと

\[A'(F) = x \exp(F(x)) G'(x \exp F(x)) - 1\]

となり \(A(F), A'(F)\) は Kinoshita-Li 合成を用いて高速に計算できるため、ニュートン法が機能します。

また、もう少し式変形を進めて Kinoshita-Li 逆関数に帰着する方法もあります。

\[H(x) = \exp G(x)\]

とおくと

\[x \exp F(x) = x H(x \exp F(x))\]

\[\frac{x \exp F(x)}{H(x \exp F(x))} = x\]

という式を得られるため、\(x \exp F(x)\) を Kinoshita-Li 逆関数で \(N+1\) 次まで計算してから形式的冪級数の \(\log\) で \(\lbrack x^N \rbrack F(x)\) を求める、という方法で解くことが出来ます。

いずれの方法でも \(\mathrm{O}(N \log^2 N)\) でこの問題を解くことが出来てこれは十分高速です。

posted:
last update: