今年中に理解する！多項式、母関数、形式的べき級数の競プロでの実践的使い方

この記事はCompetitive Programming (1) Advent Calendar 2019の7日目の記事です。

対象読者

解説で多項式とか母関数とか形式的べき級数とか書いてあるとそっ閉じするあなた
厳密な話は要求しないから、テクニックとして理解したいあなた
🤔.oO（多項式の計算についてはライブラリを使うから、それを使うまでを理解したい）

注意：以下、なんとなくで厳密性に欠ける話しかしない。概念を理解できたら幸いだ

（頑張って書いたから前半だけでも読んでって！）

第一歩「Array Restoring」

HackerRank Array Restoring

まずは、数列を母関数に変換する。
「母関数が分からんから、見てるんだけど」という声が聞こえるので説明する。
数列を母関数には以下のように変換する。

$\displaystyle{ B_1, B_2, B_3, ..., B_M \rightarrow B_1 + B_2 T + B_3 T^2 + ... + B_M T^{M-1} }$

係数に数列を置いた多項式のこと。この母関数をb(T)としておこう。
すべてが0の数列Aも母関数a(T)として表現してみると、こうだ。

$\displaystyle{ a(T) = 0 + 0 T + 0 T^2 + ... + 0 T^{N-1} }$

すると、ある操作(L,R)の計算は以下のように表現できる。

$\displaystyle{ next_a(T) = a(T) + T^{L-1} b(T) }$

b(T)を T^{L-1} 倍すると、B_1がB_1 T^{L-1}となり、B_2がB_2 T^Lのようになり、他も同様。
この状態で足すと、B_1はa(T)のT^{L-1}の係数、つまり、A_Lに足されることになる。
これで操作が多項式に対する計算に帰着されることがわかるだろう。最初、a(T)は係数が全て0なので省略されるため、Q回操作を行ったときのa(T)は以下のようになる。

$\displaystyle{ a(T) = T^{L_1-1} b(T) + T^{L_2-1} b(T) + ... + T^{L_Q-1} b(T) }$

これを式変形する。

$\displaystyle{ \begin{eqnarray} a(T) &=& T^{L_1-1} b(T) + T^{L_2-1} b(T) + ... + T^{L_Q-1} b(T) \\ &=& ( T^{L_1-1} + T^{L_2-1} + ... + T^{L_Q-1} ) b(T) \\ b(T) &=& a(T) / ( T^{L_1-1} + T^{L_2-1} + ... + T^{L_Q-1} ) \end{eqnarray} }$

この計算を行い、b(T)の係数を取得すれば、それが数列Bとして復元できるという寸法。
多項式÷多項式はO(NlogN)でやる手法がある。
この記事ではとりあえず概念理解が重要なので、そこは説明しない。
（概念理解がいまいちなうちは、実装ミスか考察ミスかわからないので実績のあるライブラリを借りることをおすすめする）
（大体そのまま使い続けちゃうんですけどね🙁）

自分の解法
 多項式の割り算はbeetさんのライブラリ、その中で使われる畳み込みはmathさんのライブラリを使っている。

int N, M, Q;
//---------------------------------------------------------------------------------------------------
void _main() {
    using T = mint;
    FormalPowerSeries<T> FPS([&](auto a, auto b) {
        MathsNTTModAny ntt;
        return ntt.solve(a, b);
    });

    cin >> N >> M >> Q;
    vector<T> fc(N);
    rep(i, 0, Q) {
        int L, R; cin >> L >> R;
        L--;
        fc[L] += 1;
    }
    vector<T> fa(N);
    rep(i, 0, N) {
        int a; cin >> a;
        fa[i] = a;
    }

    auto fb = FPS.div(fa, fc);
    rep(i, 0, M) {
        if (i) printf(" ");
        printf("%d", fb[i].get());
    }
    printf("\n");
}

母関数にすることでパワフルな操作ができることが理解できただろう。
あまり気にしなかったが、結合則や逆元をサラッと使用している。自分もあんまり気にしてない。

休憩

これで多項式の強さが分かっただろう。
だが、これはまだまだ序の口だ。
さあ、次の章に進もう。

DP更新と多項式 [yukicoderの過去問前編]

yukicoder No.3046 yukicoderの過去問

DP遷移を多項式に絡めて解くことができる。
dp[i][dan] := i回の遷移でdan段に到着する組み合わせ
これを考えると、dp[i][dan]の遷移先はdp[i+1][dan+x1],dp[i+1][dan+x2], ..., dp[i+1][xN]となる。
ここでdp[i]を数列と考えて、母関数f(T)を作成する。

$\displaystyle{ dp_{i,0} , dp_{i,1} , ..., dp_{i,K} \rightarrow f(T) = dp_{i,0} + dp_{i,1} T + ... + dp_{i,K} T^K }$

これで、dp[i][dan]は T^{dan} の係数となるわけで、遷移先を考えてみると、T^{dan+x1}, T^{dan+x2}, ..., T^{dan+xN} にこの係数をうまく分配できればいい。
ん？分配？

$\displaystyle{ dp_{i,dan} T^{dan} (T^{x1} + T^{x2} + ... + T^{xN}) = dp_{i,dan} T^{dan+x1} + dp_{i,dan} T^{dan+x2} + ... + dp_{i,dan} T^{dan+xN} }$

なるほど、(T^{x1} + T^{x2} + ... + T^{xN})をかければ1回分の遷移が行えそうだ。
g(T) = T^{x1} + T^{x2} + ... + T^{xN}としておこう。
正確には、1回分の遷移ではdp[i][0]～dp[i][K]まで行い、それぞれの組み合わせ（係数）を足し合わせる。よって、1回分の遷移を行ったdp[i+1]の母関数を得るには、以下のようにする。

$\displaystyle{ next_f(T) = f(T) g(T) }$

最初のf(T)=1であるので、i回遷移させた時のDPの母関数は以下のようになる。

$\displaystyle{ f_i(T) = {g(T)}^i }$

今回の問題では、1回以上の遷移で到達できることを考える。実際考える組み合わせは、全ての回数（1回～K回）の遷移を足し合わせたものになる。

$\displaystyle{ \sum _ {i=1}^K f_i(T) = \sum _ {i=1}^K {g(T)}^i }$

この計算によって出てきた多項式の T^K の係数が求めるべき答えになっている。
しかし、このままではTLE。
数式続きで疲れてきたかもしれないが、ここから、また新しい概念を導入しよう。

そして、形式的べき級数へ [yukicoderの過去問後編]

形式的べき級数とは、無限項の多項式のことである。
無限にすることでできることがある。
今回でいうと、以下の公式を適用できる（どんな場合でも適用できるかは知らない）。

$\displaystyle{ \sum _ {i=1}^{\infty} {f(T)}^i = \frac{1}{1-f(T)} }$

前章ではK回までしか考えなかったが \infly 回までとしても答えは変わらない。
よって、

$\displaystyle{ \frac{1}{1-g(T)} }$

これの T^K の係数が答えとなる。
やっとスッキリした形に落ちた。
多項式÷多項式はO(NlogN)でできるので、この問題も解けた。
とりあえず「形式的べき級数を導入すると嬉しくなる」パターンを1つ分かってもらえたかと思う。

自分の解法
 多項式の割り算はbeetさんのライブラリ、その中で使われる畳み込みはmathさんのライブラリを使っている。

int K, N, x[101010];
//---------------------------------------------------------------------------------------------------
void _main() {
    using T = mint;
    FormalPowerSeries<T> FPS([&](auto a, auto b) {
        MathsNTTModAny ntt;
        return ntt.solve(a, b);
    });

    cin >> K >> N;
    rep(i, 0, N) cin >> x[i];

    vector<T> fy(K + 1);
    fy[0] = 1;
    rep(i, 0, N) fy[x[i]] = -1;
    fy = FPS.inv(fy, K + 1);

    cout << fy[K] << endl;
}

休憩と次に繋がる知識

ここまで来ることができれば、だいぶ感覚は掴めてきている。
次は実践編だと行きたいが、次につながる重要な考えを前の解法から抽出しておきたい。
前の問題で、

$\displaystyle{ next_f(T) = f(T) g(T) }$

この式が出てきたが、母関数としては多項式の積であるが、元々の数列的には畳み込み計算を行っていることになる。
畳み込み計算とは「C[i]がa+b=iを満たすA[a]*B[b]の総和」のことである。
くどいかもしれないが確認すると「畳み込み計算は母関数にすると関数の積になる」ということである。
さて、おまたせしてすまない。次の章に行こうか。

木DPで母関数？？？ [The Child and Binary Tree 前編]

Codeforces The Child and Binary Tree
問題概要

以下の条件を満たす根付き二分木を構築する組み合わせは？
- 頂点数はなんでもいい
- 各頂点に割り当てられる数はc1～cNのどれか（複数回用いても良い）
- 割り当てた数の総和がiになる
この質問を[1,M]で全てのiについて答える（mod 998244353)

従来の勘をたどると、まずは木DPっぽことをする感じに見える。
dp[tot] = 数の合計がtotとなる木の組み合わせ
この漸化式を考える。
根にrootを選択して、左に総和left、右に総和rightをくっつける組み合わせを考えると、
dp[root + left + right] += dp[left] * dp[right]
となる。
DPの常套テクとして、このような配るDPは貰うDPにすると最適化しやすいのでやってみる。

$\displaystyle{ dp_{tot} = \sum _ {root_i=1}^{N} \sum _ {left=0}^{tot-C_{root\_i}} dp_{left} * dp_{tot-C_{root\_i}-left} }$

このような二重和が出てくる。
厄介な二重和だ。しかも添字になっているというのも使いにくい。
根として使う数も数として扱うために「c_cmb[c] := 数cが使えれば1。そうでないなら0」を導入する。
すると、これは組み合わせっぽく使えるので、以下のように書き直すことができる。

$\displaystyle{ dp_{tot} = \sum _ {root=0}^{tot} \sum _ {left=0}^{tot-root} c_com_{root} * dp_{left} * dp_{tot-root-left} }$

勘が良い方は見えているかもしれないが、これは畳込み計算である。
「dp[tot]がroot+left+right=totを満たすc_com[root]dp[left]dp[right]の総和」になっている。
前の章で確認した「畳み込み計算は母関数にすると関数の積になる」を使う時が来た。
母関数にしよう。

母関数を解く [The Child and Binary Tree 後編]

dpを母関数f(T)とする。先程定義したc_comも母関数c(T)とする。

$\displaystyle{ dp_0, dp_1, ..., dp_M \rightarrow f(T) = dp_0 + dp_1 T + ... + dp_M T^M \\ c\_com_0, c\_com_1, ..., c\_com_M \rightarrow c(T) = c\_com_0 + c\_com_1 T + ... + c\_com_M T^M \\ }$

すると、先程の畳み込み計算式は以下のようになる。

$\displaystyle{ f(T) = c(T) {f(T)}^2 + 1 }$

さて、行間がかなりあるので、一応下に箇条書きで説明書きというかヒントを書いておく。
根性がいるし、もしかしたら脱落するかもしれない。
ここが踏ん張り所なので、頑張ってみてほしい。

c(T)とf(T)とf(T)の畳み込み計算なので、c(T) {f(T)}^2 になっている
畳み込みの計算の結果がdpなので、左辺にはf(T)が入る
右辺の1は何か
- これはdpの初期値1である
- dp[0]はc_com[0]dp[0]dp[0]からしか計算されないが、c_com[0]は0であるため、c(T) {f(T)}^2 だけだと定数項がなくなってしまう。よって、dp[0]=1として1を足している
- f(T) = x0 + x1 T + x2 T² + ...となると思うが、x0=1となる
- 問題から見ると数の総和0の根付き二分木を作る組み合わせは1
f(T)が左辺右辺のどちらにもあるが、これは等式か代入式か
- これは等式である
- これから先はちゃんと勉強してなくてイメージで解釈したことを説明する（間に受けない方が良いかも）
- 1. 例えば右辺でc_com[3]dp[2]dp[2]とかが計算されて、dp[7]が更新される
- 1. 等式なので、それは左辺のdp[7]がその値になっている
- 1. 等式なので、それは右辺のdp[7]もその値になっている
- 1. dp[7]は更新されたが、この更新によって、再度dp[7]が書き換わることはない（c_com[0]=0なので）
- 1. dp[7]が確定することで、dp[8]が計算可能になり、順番に計算が進んでいく
- 1. 無限に計算が進み左辺と右辺が等しくなることはなさそうだが、形式的べき級数なので、無限に発散した（収束？）やつはなんかうまいこと消え去って等しくなる（1=0.99999...が成り立つ的なイメージ）

競プロドリブンで勉強したので、この説明を教授にして怒られても責任は取れない。
さて、ここまでおぼろげに分かっていれば峠は超えたので、答えに行こう。
先程の式はf(T)に注目すると2次方程式なので、中学で習う解の公式を使って答えを求めよう。

$\displaystyle{ f(T) = \frac{1 ± \sqrt{1 - 4c(T)}}{2 c(T)} }$

いつものここを参考にすると、「初期条件 f(0)=c0=1 であり，x=0 で関数が連続になる必要があるため符号がマイナスの方が採用される。」と書いてある。
この問題の公式解説もマイナスを採用しているので、そういうものなんだろう。

$\displaystyle{ f(T) = \frac{1 - \sqrt{1 - 4c(T)}}{2 c(T)} }$

これで今までやってきた「多項式計算」に帰着できた。おめでとう。
だが、このまま計算すると、分母の定数項が0になってしまい、計算できない（なんかね、invがとれないっぽい）
なので、有理化っぽいことをやって、2つあるc(T)を1つにして、分母の定数項が0にならないようにしよう。
式変形はtexを書くのが面倒なので、公式解説の「438E - The Child and Binary Tree」の部分に同じ式があるので見てみると良い。
英語は読まなくていい。俺が日本語でここに書いてある。

$\displaystyle{ f(T) = \frac{2}{1 + \sqrt{1 - 4c(T)}} }$

後は実装するだけだ。
多項式の平方根が出てくるがO(NlogN)でできる。

自分の解法
 多項式の割り算と平方根計算はbeetさんのライブラリ、その中で使われる畳み込みはmathさんのライブラリを使っている。

const int MAXSIZE = 101010;
int N, M, C[101010];
MathsNTTModAny ntt;
//---------------------------------------------------------------------------------------------------
void _main() {
    using T = mint;
    FormalPowerSeries<T> FPS([&](auto a, auto b) {
        MathsNTTModAny ntt;
        return ntt.solve(a, b, 998244353);
        });
 
    cin >> N >> M;
    rep(i, 0, N) cin >> C[i];
    
    vector<T> sqr(MAXSIZE, 0);
    sqr[0] = T(1);
    rep(i, 0, N) sqr[C[i]] -= T(4);
    sqr = FPS.super_sqrt(sqr, MAXSIZE);
 
    vector<T> dwn = sqr;
    dwn[0] += 1;
 
    vector<T> up(MAXSIZE, 0);
    up[0] = 2;
 
    vector<T> ans = FPS.mul(up, FPS.inv(dwn, MAXSIZE));
    for (int i = 1; i <= M; i++) printf("%d\n", ans[i].get());
}

終わりに

長文とこのノリに付き合ってくれてありがとう。
個人的にかなりいい記事に仕上がっていると思うし、今年はちょっとバズったこの記事も輩出しているので、
競プロブログ・オブ・ザ・イヤーは「はまやんはまやんはまやん」に違いない。
トロフィーの到着を待ちながら年を越すことにしよう。

免責事項と本当の終わりに

多分数学的にはガバガバ記事なので、そこら辺気にする人は裏付けを取って使ってください。
あえて厳密性を避けた所もありますが、多分だいたい自分の理解不足だと思います。
ちゃんと体系的に勉強しましょうということですね。
つよつよ数学家からクリスマスプレゼントが届かないことを祈りつつ、締めとします。
2020年もよろしくお願いいたします。