形式的冪級数

概要

多項式の展開計算ルールを用いて、複雑な組み合わせの計算等を行う方法。
英語のFormal Power Seriesの頭文字を取ってFPSと呼ばれることもある。

かなり数学よりの話で、手計算の数学でも用いることがある。
競技プログラミングにおいても、出題されるとそれなりに手計算が必要となることも多い。
このページの解説でも、ただの数学の解説をしているだけの項目も多い。

畳み込み、高速フーリエ変換（FFT）、数論変換（NTT）もこのページで解説する。

アルゴリズム内容

例えば、以下のような問題を考える。

6面体のサイコロを20回振る。
出た目の合計がkになる目の出方を、k=20,21,22,……,120についてそれぞれ求めよ。

$k=20$ から $k=25$ までは、重複組み合わせとして解ける。
$k=26$ になると、サイコロに $7$ の目がない分を考慮しなくてはならない。
$k=32$ から先は、$7$ 以上の目が出る組み合わせを引くときに、ダブって引かないように注意が必要。
ちょうど真ん中の $k=70$ などは、もはや何から手を付けていいかわからないほどである。

こんな感じの複雑な組み合わせを求めるのに役に立つのが、形式的冪級数である。

形式的冪級数

なんらかの値を係数に持つ、$x$ の無限次元多項式の形をしている式のこと。
途中から係数がすべて $0$ であることにすれば、有限次元多項式もこの中に含まれる。

$$
f(x) = a_0x^0 + a_1x^1 + \dots + a_kx^k + \dots
$$

無限次元多項式「の形」というのは、式としての意味や正当性を気にしないということ。

例えば、数学の文章題では $x$ という文字になんらかの意味を持たせることが多い。
しかし、形式的冪級数ではそれは気にしない。
多項式や展開公式の形を利用したいだけで、$x$ の中身が何なのかなど最初から考える気がないのである。
（逆に、数学でのただの「次の式を展開せよ」という問題で扱う式は形式的冪級数と言えるかもしれない）

また、以下の冪級数は、形式的でない場合は $-1<x<1$ という制約をつけなければ定義できない。
なぜなら、その範囲外の値を代入すると、無限級数が発散してしまうから。

$$
f(x) = x^0 + x^1 + \dots + x^k + \dots
$$

しかし、形式的冪級数ではそんなことは全く気にしない。
多項式や展開公式の形を利用したいだけで、$x$ の中身が何なのかなど最初から考える気がないのである。

通常型母関数

数列の表し方について考える。
例えば、初項が $A_0=3$ で、公比 $2$ である等比数列 ${A_n}$ を表す方法をいろいろ考えてみる。

まず、最も単純な方法として、書き並べるというものがある。
$$
{A_n} = 3, 6, 12, 24, 48, \dots
$$

それから、一般項を書くという方法がある。
$$
A_n = 3\times 2^n
$$

さらに、漸化式を使うという方法もある。
$$
A_0 = 3, \quad A_{n+1} = 2A_n \quad (n=0,1,2,\dots)
$$

ここまでは高校数学で取り扱う内容である。
そして、第 $4$ の表し方として、（通常型）母関数という表し方を導入する。
$$
A(x) = \dfrac{3}{1-2x}
$$

一見先ほどまでの $A_n$ とは全く無関係な式に見える。
が、そのつながりは、形式的にマクローリン展開してみればわかる。

$$
A(x) = \dfrac{3}{1-2x} = 3x^0 + 6x^1 + 12x^2 + 24x^3 + 48x^4 + \dots
$$

見てわかる通り、$x^k$ の係数が $A_k$ となっているのである。
このように、数列の値を係数とする形式的冪級数で数列を表す方法を、（通常型）母関数という。
（通常型というのは、後に特殊な係数の定め方をする場合があるためで、今は深く考えなくてよい）
ちなみに、名前に関数と入っているが、$x$ に値を入れることなど考えていないので、実は関数ではない。

この等比数列の場合は母関数が数列とは直感的に結びつかない形であったが、もっと単純な場合もある。
特に有限数列の場合は、ほぼ見た目通りの母関数になる。
例えば、$3,1,4,1,5$ という長さ $5$ の数列の母関数は、$A(x) = 3x^0+1x^1+4x^2+1x^3+5x^4$ である。
競技プログラミングでは、真の意味での無限次元は扱えないので、このようなものの方が扱う機会は多い。

母関数は、一般項や漸化式から求めることができる。
特に、恒等数列や等比数列の場合は一般項から求めるとよい。
初項が $A_0=3$ で、公比 $2$ である等比数列 ${A_n}$ の場合は、まず愚直な形式的冪級数が以下のようになる。
$$
A(x) = 3x^0 + 6x^1 + 12x^2 + 24x^3 + 48x^4 + \dots
$$

これを、初項が $3$ で公比 $2x$ の無限等比級数だと思うことにする。
そして、形式的に無限等比級数の公式を用いることで、以下が得られる。
$$
A(x) = \dfrac{3}{1-2x}
$$

漸化式の方は、例えば、初項が $a_0=2$ であり $a_{n+1} = 2a_n-1$ で定められる数列の母関数を求めてみる。
まず、この漸化式の両辺に $x^{n+1}$ をかけて、$n\geq 0$ 範囲で総和を取る。
$$
\sum_{n=0}^{\infty} a_{n+1}x^{n+1} = 2x \sum_{n=0}^{\infty}a_nx^n - \sum_{n=0}^{\infty} x^{n+1}
$$

$\sum_{n=0}^{\infty}a_nx^n$ は母関数そのもので、$\sum_{n=0}^{\infty} a_{n+1}x^{n+1}$ はそこから $x^0$ 項が消えたもの。
また、$\sum_{n=0}^{\infty} x^{n+1}$ はただの無限等比級数。
よって、これは以下のように変形できる。
$$
A(x) - 2 = 2x A(x) - \dfrac{x}{1-x}
$$
これを $A(x)$ について解けば、母関数が得られる。
$$
A(x) = \dfrac{2-3x}{(1-x)(1-2x)}
$$

練習として、フィボナッチ数列（$F_0=0,F_1=1,F_{n+2}=F_{n+1}+F_n$）の母関数を求めてみるとよい。
結果は $F(x) = \dfrac{x}{1-x-x^2}$ になるはずである。

また、非負整数定数 $m$ に対する $A_n={}_{n+m}\mathrm{C}_{m}$ の母関数 $A_m(x)$ を求めて覚えておくと、非常に役に立つ。
これも一度自力で作ってみるとよい。
（ヒント：$A_n=1$ の母関数を作って、両辺を何度も微分する）

分数型の母関数の級数展開

$r$ は正の整数として、$\dfrac{1}{(1-x)^r}$ の展開は、以下である。
これは、ここまでを丁寧に読んできた人には既知のはずである。
$$
\dfrac{1}{(1-x)^r} = {}_{r-1}\mathrm{C}_{r-1}x^0+{}_{r}\mathrm{C}_{r-1}x^1+{}_{r+1}\mathrm{C}_{r-1}x^2+{}_{r+2}\mathrm{C}_{r-1}x^3+\dots
$$

$k\neq 0$ である場合の $\dfrac{1}{(k-mx)^r}$ も、$\dfrac{1}{k^r}$ を括りだしてから上の公式を使えば、以下のようになる。
$$
\dfrac{1}{(k-mx)^r}
=
\dfrac{1}{k^r}
\left\{
{}_{r-1}\mathrm{C}_{r-1}x^0
+
{}_{r}\mathrm{C}_{r-1}\left(\dfrac{m}{k}\right)x^1
+
{}_{r+1}\mathrm{C}_{r-1}\left(\dfrac{m}{k}\right)^2x^2
+
{}_{r+2}\mathrm{C}_{r-1}\left(\dfrac{m}{k}\right)^3x^3
+\dots
\right\}
$$

これで、分母が $1$ 次式の累乗である場合は、正確に無限級数化できる。
分子が $1$ でなかったとしても、分子を $1$ に変えた式の級数展開に分子を掛け算すればよい。

問題は、分母が $2$ 次以上の式だった場合。
これも分子が $1$ のものだけ処理できればよいが、しかしそれでも正確に一般項を表現できるとは限らない。
とはいえ、競技プログラミング的には十分な次数まで求まっていれば実質的に無限に求めたとみなせる。
その妥協が許されるのならば、これはNewton法を用いることができる。

$f(0)\neq 0$ である $f(x)$ の展開は既知であるものとして、$g(x)=\dfrac{1}{f(x)}$ の展開を考える。
$x^k$ の項まで展開したもの $g_k(x)$ とする。

まず、$g_0(x)$ を求める。
これは、 $f(x)g(x)=1$ であることから、$g(0)=\dfrac{1}{f(0)}$ すなわち $g_0(x)=\dfrac{1}{f(0)}$

これを元手にして、次数の正確性を上げていく。
$g_k(x)$ が求まったとしよう。
これは真の $g(x)$ を $x^k$ まで正確に表現しているので、$e(x)=g(x)-g_k(x)$ は $O(x^{k+1})$ である。
このとき、$h(x) = g_k(x)\{2-f(x)g_k(x)\}$ と定義し、変形してみる。
$$
\begin{aligned}
h(x)
&=g_k(x)\{2-f(x)g_k(x)\}\\
&=\{g(x)-e(x)\}\{2-f(x)g(x)+f(x)e(x)\}\\
&=\{g(x)-e(x)\}\{1+f(x)e(x)\}\\
&=g(x)-e(x)+f(x)g(x)e(x)-f(x)\{e(x)\}^2\\
&=g(x)-f(x)\{e(x)\}^2
\end{aligned}
$$

$e(x)=g(x)-g_k(x)$ は $O(x^{k+1})$ なので $f(x)\{e(x)\}^2$ は $O(x^{2k+2})$ である。
したがって、$h(x)$ は $x^{2k+1}$ まで正確であるため、そこから先を切り落とせば $g_{2k+1}(x)$ となる。
これを繰り返せば、$g_0(x)\to g_1(x)\to g_3(x)\to g_7(x)\to g_{15}(x)\to \dots$ と求めることができる。
$g_k(x)$ を求める場合、これを $\log_2(k+1)$ 回以上繰り返せばよい。

一例として、フィボナッチ数列の母関数 $F(x) = \dfrac{x}{1-x-x^2}$ を級数展開してみよう。
$f(x) = 1-x-x^2,g(x) = \dfrac{1}{1-x-x^2}$ とする。

$f(0) = 1$ より、$g_0(x) = 1$ である。
$g_0(x)\{2-f(x)g_0(x)\} = 1+x+x^2$ を $x^1$ 項まで採用して、$g_1(x) = 1+x$ となる。
$g_1(x)\{2-f(x)g_1(x)\} = 1+x+2x^2+3x^3+x^4$ から、$g_3(x) = 1+x+2x^2+3x^3$ となる。
同様にすると $g_7(x) = 1+x+2x^2+3x^3+5x^4+8x^5+13x^6+21x^7$ となる。

分子の $x$ も掛けると、$\dfrac{x}{1-x-x^2} = x+x^2+2x^3+3x^4+5x^5+8x^6+13x^7+21x^8+\dots$ となる。
このように、実際にフィボナッチ数列が再現される。

プログラムで行う場合の計算量は基本的に級数同士の積を求める計算量の $\log k$ 倍である。
後述のFFT/NTTを用いた場合、$1$ 回の積は $O(k\log k)$ なので、全体の計算量は $O(k\log^2 k)$ である。

ただし、$g_3(x)$ を求めるときに $f(x)$ も $x^3$ まであれば十分、などの工夫をすれば少し早くなる。
FFT/NTTを用いた上にその工夫もすれば、計算量は $O(k\log k)$ になる。

畳み込み

$2$ つの数列に対して、「畳み込み」と呼ばれる演算処理が存在する。
$A$ と $B$ の添字の和が $k$ になる組全ての積を求め合計し、それを $C_k$ とする。
$$
C_k = A_0B_k + A_1B_{k-1} + A_2B_{k-2} + \dots + A_kB_0
$$

有限数列の末項より後ろは $0$ とみなしながらこの $C_k$ を全て求める計算が、畳み込みである。
（ある $k$ での値 $1$ つだけ求めるだけのものも畳み込みと呼ぶことがある）

例えば、$\{A_n\} = 3,2,1$ で $\{B_n\} = 20,15,10$ だとする。
この場合、$C_2 = 3\times 10+2\times 15+1\times 20 = 80$ となる。
同様に、全ての $k$ の値で求めて数列化すると、$\{C_n\} = 60,85,80,35,10$ となる。

画像処理などでよく用いられるらしいが、競技プログラミングで画像処理などしない。
そのため、これが何の役に立つのかは、後の「組み合わせ問題への適用」の項目を参照。

これだけ聞くととてもややこしい計算に見える。
しかし、実は形式的冪級数にすると、その意味は非常に明快である。
$(3x^0+2x+1x^2)\times(20x^0+15x^1+10x^2) = 60x^0+85x^1+80x^2+35x^3+10x^4$ ということである。

愚直に計算するのは簡単なコードで実装でき、計算量は $O(|A||B|)$ である。
次に述べるFFT/NTTを用いれば $O((|A|+|B|)\log(|A|+|B|))$ に高速化できる。

高速フーリエ変換（FFT）

長さ $|A|$ の数列 $A$ と長さ $|B|$ の数列 $B$ の畳み込みを愚直に求めると、$O(|A||B|)$ かかる。
しかし、高速フーリエ変換（FFT）を利用すると $O((|A|+|B|)\log(|A|+|B|))$ で畳み込みを行える。

ここでは、その原理の数学的説明を述べる。
が、実はこの点については、AtCoder Libraryのconvolutionを利用することも可能。
「中身わからなくてもライブラリ使って問題解ければ十分です！」の場合は、この項目はスキップ推奨。

$A(x)=2+x$ と $B(x)=1-2x$ の畳み込み $C(x)$ を、変換を用いて求める方法を考える。

まず、普通の複素離散フーリエ変換による畳み込みを考えてみる。
普通のというのは、計算量が愚直な畳み込みから改善していない（むしろ悪化する）という意味である。

まず、畳み込み後の次数を考えると $1+1=2$ 次となるので、それより大きい数として $4$ を選ぶ。
つまり、畳み込みを $4-1=3$ 次式として求める。
（$3$ を選ぶと以後の計算が少し面倒になり、しかも高速化は同じ方法ではできない）

まず、$4$ 乗すれば $1$ になるが、$3$ 乗以下では $1$ にならない複素数を $1$ つ選ぶ。
今回は、虚数単位 $i$ を選ぶことにする。
$A$ と $B$ それぞれ、$i^0=1,i^1=i,i^2=-1,i^3=-i$ を代入して値を求め、掛けて $C$ での値にする。

• $A(1) = 3, B(1) = -1$ より $C(1) = -3$
• $A(i) = 2+i, B(i) = 1-2i$ より $C(i) = 4-3i$
• $A(-1) = 1, B(-1) = 3$ より $C(-1) = 3$
• $A(-i) = 2-i, B(-i) = 1+2i$ より $C(-i) = 4+3i$

これを逆変換にかければよい。
つまり、以下の式に代入して計算する。
$$
c_j = 4^{-1} \sum_{k=0}^{3} C(i^k) i^{-jk}
$$

すると、以下のようになる。

• $c_0 = 4^{-1}\times \{ -3\times i^0+(4-3i)\times i^0+3\times i^0+(4+3i)\times i^0\} = 2$
• $c_1 = 4^{-1}\times \{ -3\times i^0+(4-3i)\times i^{-1}+3\times i^{-2}+(4+3i)\times i^{-3}\} = -3$
• $c_2 = 4^{-1}\times \{ -3\times i^0+(4-3i)\times i^{-2}+3\times i^{-4}+(4+3i)\times i^{-6}\} = -2$
• $c_3 = 4^{-1}\times \{ -3\times i^0+(4-3i)\times i^{-3}+3\times i^{-6}+(4+3i)\times i^{-9}\} = 0$

よって、$C(x)=2-3x-2x^2$ である。

念のため、一般の関数でこれが可能な証明の概略も示しておく。
$1$ の原始 $N$ 乗根をとり、$\xi$ とする。
そして、$C(\xi) = A(\xi)B(\xi)$ を求めるまでは、ただの作業しかしていない。
最後に $c_n$ を取り出すところだが、
$$
N^{-1} \sum_{k=0}^{N-1} C(\xi^k) \xi^{-jk}
= N^{-1} \sum_{k=0}^{N-1} \sum_{l=0}^{N-1} c_l \xi^{lk} \xi^{-jk}
= N^{-1} \sum_{k=0}^{N-1} \sum_{l=0}^{N-1} c_l \xi^{(l-j)k}
= N^{-1} \sum_{l=0}^{N-1} c_l N\delta_{jl}
= c_j
$$
となる。
ただし、$\delta_{ij}$ はクロネッカーのデルタ。

以上が、複素離散フーリエ変換である。

既に述べた通り、このままでは計算量が $O(|A||B|)$ どころか $O(N^2)$ で悪化している。
しかし、最初に選ぶ数 $N$ が $2$ の累乗である場合、これを高速化できる。

関数 $A$ を、偶数次と奇数次で分ける。
$A(x) = A_E(x^2) + xA_O(x^2)$
$A_E(\xi^{2k})$ と $A_O(\xi^{2k})$ の値がわかれば、$A(\xi^k)$ は $1$ つあたり $O(1)$ で求められる。
そして、$2$ つの関数 $A_E(x)$ と $A_O(x)$ は $A(x)$ と比べて長さが半分になっている。
しかも、$\xi$ が $1$ の原始 $N$ 乗根であれば、$\xi^2$ は $1$ の原始 $N/2$ 乗根である。

ということは、これを再帰的に行うことができ、$\log_2{N}$ 回行うと、長さ $1$ になる。
長さ $1$ の複素離散フーリエ変換は明らかにそのまま何もしないことである。
よって、$O(N\log N)$ で変換はできた。

次に逆変換だが、実はよく見ると普通の変換とほぼ同じことをしている。
違いは以下の $3$ 点。

• $x^k$ 係数に $C(\xi^k)$ を用いている
• $x$ に入れる値に $\xi^0,\xi^(-1),\xi^(-2),\dots,\xi^(-N+1)$ を用いている
• 最後に $N$ で割る処理がある

これは高速化処理のロジックに影響を与えない。
よって、逆変換も $O(N\log N)$ でできた。

間の $N$ 個の掛け算は当然 $O(N)$ で終わるので、これで全体が $O(N\log N)$ となった。
$N$ として、$(|A|+|B|-1)$ 以上の最小の $2$ の累乗を選べば $O((|A|+|B|)\log(|A|+|B|))$ である。
これが高速フーリエ変換（fast Fourier transform）である。

ただし、これを機械計算で行うのは、計算量とは別の問題がある。
一般に $N\geq3$ の場合、$1$ の原始 $N$ 乗根は、非整数である。
そのため、小数計算で途中の処理を行い、最後に整数にまとめるという強引な方法を取らざるを得ない。
この問題を解決するのが、次に述べる数論変換である。

数論変換（NTT）

数論変換の原理の数学的説明を述べ、コード例の項目で実装方法を解説する。
「中身わからなくてもライブラリ使って問題解ければ十分です！」の場合は、この項目はスキップ推奨。

離散フーリエ変換には、$N\geq3$ の場合 $1$ の原始 $N$ 乗根が整数にならないという問題点があった。
しかし、ある素数 $p$ を法とした有限体 $\mathbb{F}_p$ 上ではどうだろう。
（有限体については、$p$ が素数なので割り算も許される剰余類環、くらいに理解していれば十分）

例えば有限体 $\mathbb{F}_5$ において、$2^1=2,2^2=4,2^3=3,2^4=1$ であるため、$2$ は $1$ の原始 $4$ 乗根である。
つまり、先ほどの畳み込みは、$i$ ではなく $2$ を用いて以下のように計算できる。
全て $\bmod 5$ で計算していることに注意。

• $A(x)=2+x$
• $B(x)=1-2x=1+3x$

• $A(1) = 3, B(1) = 4$ より $C(1) = 2$
• $A(2) = 4, B(2) = 2$ より $C(2) = 3$
• $A(4) = 1, B(4) = 3$ より $C(4) = 3$
• $A(3) = 0, B(3) = 0$ より $C(3) = 0$

• $c_0 = 4^{-1}\times \{ 2\times 2^0+3\times 2^0+3\times 2^0+0\times 2^0\} = 2$
• $c_1 = 4^{-1}\times \{ 2\times 2^0+3\times 2^{-1}+3\times 2^{-2}+0\times 2^{-3}\} = 2$
• $c_2 = 4^{-1}\times \{ 2\times 2^0+3\times 2^{-2}+3\times 2^{-4}+0\times 2^{-6}\} = 3$
• $c_3 = 4^{-1}\times \{ 2\times 2^0+3\times 2^{-3}+3\times 2^{-6}+0\times 2^{-9}\} = 0$

よって、$C(x)=2+2x+3x^2$ となる。
本来の計算結果である $C(x)=2-3x-2x^2$ と、$\bmod 5$ で一致している。
このようにして、$\bmod p$ での答えでよければ整数計算だけで答えが出せる。

しかし、これを実用する上では $1$ つ問題がある。
高速変換のためには、$N$ の値として $(|A|+|B|-1)$ 以上の最小の $2$ の累乗を選ぶ必要がある。
しかし、先に $N$ を決めてから $\mathbb{F}_p$ 上での原始 $N$ 乗根を見つけるのは非常に難しい。
最悪の場合、そもそも存在しない場合すらある。

そこで競プロでよく用いられるのが、$\mathbb{F}_{998244353}$ という有限体。
$\mathbb{F}_{998244353}$ 上では、実は $3$ が原始 $119\times2^{23}$ 乗根になっているのである。
ということは、$3^{119\times2^{23-k}}$ という数を用意すれば、これは $1$ の原始 $2^k$ 乗根。
$k$ の値として最大で $23$ まで選べるため、長さも十分。
これで足りないようなら、そもそも配列の長さが $10^7$ くらいあるということ。
つまり、$O(N\log N)$ がちゃんと走ったところでどうせ間に合わないのである。

ということで、$\mod{998244353}$ で割った余りを答える問題でなら、実用的にNTTを使える。
$p-1$ 乗などの計算が入るため、計算量は $O((|A|+|B|)\log(|A|+|B|))+\log p\log(|A|+|B|)$ となる。
（工夫すれば $O((|A|+|B|)\log(|A|+|B|))+\log p$ にもできるが、労力に見合わない）

$998244353$ 以外の法では、原始 $2^k$ 乗根を気合で探す分が上乗せになる。
（Google検索等で原始 $2^k$ 乗根の情報が手に入れば、その法でも使える）

余談だが、何かで割って余りを答える問題で、ほとんどの場合に法が $998244353$ なのはここが原因。
想定解がNTTのときだけ $998244353$ で割った余りを要求すると、解法がNTTだとバレるのである。

コード例

NTTの実装例を載せる。
FFTは省略。
「中身わからなくてもライブラリ使って問題解ければ十分です！」の場合は、この項目はスキップ推奨。

AtCoder Libraryのものとの違いは、以下。

• $\bmod 998244353$ 固定で、他の法でやることは考えていない
• 第 $3$ 引数 lim として、$x$ の最高次数を指定して切り落とせるようになっている
  • 「分数型の母関数の級数展開」のところで述べた、$f(x)$ を一時的に切り落とす処理も行う
  • 指定の省略も可

power_mod(a,m) という、$a^m \mod 998244353$ を求めてくれる関数が別に必要。
しかも、それが $m<0$ にも対応されている前提。
また、全ての要素が事前に $0$ 以上 $998244353$ 未満にしてある前提である。

// bit反転位置でswapして整列しなおす関数
// 偶数次と奇数次にわけていく、再帰を潜る処理に相当
// 例えば、5bitで11(01011)番目の要素と26(11010)番目の要素を入れ替える
void bit_rev(vector<long long>& a, int len) {
  assert(0<=len&&len<=30);
  assert(ssize(a)==(1<<len));
  if (len==0) return;
  for (int i=1, pos=0; i<(1<<len); i++) {
    int bit=(1<<(len-1));
    while (pos&bit) {pos ^= bit; bit >>= 1;}
    pos ^= bit;
    if (i<pos) swap(a[i],a[pos]);
  }
}

// バタフライ演算
// 偶数次と奇数次の結果から本来の値を求める、再帰を戻る処理に相当
// xが原始N乗根であるとき、x^(N/2)は-1であることを利用して少し効率化している
void butterfly(vector<long long>& a, int len, bool inv = false) {
  assert(mod==998244353);
  assert(0<=len&&len<=23);
  assert(ssize(a)==(1<<len));
  for (int w=1; w<(1<<len); w*=2) {
    long long x;
    if (!inv) x = power_mod(3,(mod-1)/(2*w));
    else x = power_mod(3,-(mod-1)/(2*w));
    for (int i=0; i<(1<<len); i+=2*w) {
      long long xk = 1;
      for (int j=0; j<w; j++) {
        long long ae = a[i+j];
        long long xao = xk*a[i+j+w]%mod;
        a[i+j] = (ae+xao)%mod;
        a[i+j+w] = (ae-xao+mod)%mod;
        xk = xk*x%mod;
      }
    }
  }
}

// NTTによる順変換および逆変換
void ntt(vector<long long>& a, int len, bool inv = false) {
  assert(mod==998244353);
  assert(0<=len&&len<=23);
  assert(ssize(a)==(1<<len));
  bit_rev(a,len);
  butterfly(a,len,inv);
  if (inv) {
    long long x = power_mod(1<<len,-1);
    for (long long& i : a) i = i*x%mod;
  }
}

// 畳み込み
vector<long long> convolution(const vector<long long>& a_org, const vector<long long>& b_org, int lim = (1<<23)-1) {
  if (a_org.empty()||b_org.empty()||lim<0) return {};
  vector<long long> a(a_org.begin(),a_org.begin()+min((int)ssize(a_org),lim+1));
  vector<long long> b(b_org.begin(),b_org.begin()+min((int)ssize(b_org),lim+1));
  int len = 0;
  int max_deg = ssize(a)+ssize(b)-2;
  while ((1<<len)<=max_deg) len++;
  assert(len<=23);
  a.resize(1<<len,0);
  ntt(a,len);
  b.resize(1<<len,0);
  ntt(b,len);
  for (int i=0; i<(1<<len); i++) a[i] = a[i]*b[i]%mod;
  ntt(a,len,true);
  a.resize(min(max_deg,lim)+1);
  return a;
}

// 逆数
// 途中計算では、全て符号反転状態で保持している
vector<long long> fps_inv(const vector<long long>& a, int lim) {
  assert(mod==998244353);
  assert(lim>=0);
  assert(!a.empty());
  assert(a[0]!=0);
  vector <long long> inv = {mod-power_mod(a[0],-1)};
  for (int n=1; n<=lim; n*=2) {
    vector<long long> tmp = convolution(a,inv,min(2*n-1,lim));
    tmp[0] = (tmp[0]+2)%mod;
    inv = convolution(tmp,inv,min(2*n-1,lim));
  }
  inv.resize(lim+1);
  for (long long& i : inv) if (i>0) i = mod-i;
  return inv;
}

おまけ。（むしろ本体説もある）
AtCoder Library を部分的に取り入れた逆数計算のコード。

power_mod(a,m) という、$a^m \mod 998244353$ を求めてくれる関数が別に必要。
しかも、それが $m<0$ にも対応されている前提。
また、全ての要素が事前に $0$ 以上 $998244353$ 未満にしてある前提である。

うっかり上のコードと両方並べると convolution が衝突するので注意。

// AtCoder Libraryのconvolutionのラッパー
vector<long long> convolution_acl(const vector<long long>& a_org, const vector<long long>& b_org, int lim = (1<<23)-1) {
  if (a_org.empty()||b_org.empty()||lim<0) return {};
  vector<long long> a(a_org.begin(),a_org.begin()+min((int)ssize(a_org),lim+1));
  vector<long long> b(b_org.begin(),b_org.begin()+min((int)ssize(b_org),lim+1));
  vector<long long> result = convolution(a,b);
  if (ssize(result)>lim+1) result.resize(lim+1);
  return result;
}

// 逆数
// 途中計算では、全て符号反転状態で保持している
vector<long long> fps_inv(const vector<long long>& a, int lim) {
  assert(mod==998244353);
  assert(lim>=0);
  assert(!a.empty());
  assert(a[0]!=0);
  vector <long long> inv = {mod-power_mod(a[0],-1)};
  for (int n=1; n<=lim; n*=2) {
    vector<long long> tmp = convolution_acl(a,inv,min(2*n-1,lim));
    tmp[0] = (tmp[0]+2)%mod;
    inv = convolution_acl(tmp,inv,min(2*n-1,lim));
  }
  inv.resize(lim+1);
  for (long long& i : inv) if (i>0) i = mod-i;
  return inv;
}

使い方の応用

組み合わせ問題への適用

例えば、以下の問題を考える。

AとBの2つの箱があり、それぞれ球が10個入っている。
Aには0点の球が5個、1点の球が3個、2点の球が2個入っている。
Bには0点の球が6個、1点の球が2個、2点の球が1個、3点の球が1個入っている。
k=0,1,2,…,5それぞれについて、以下の問いに答えよ。
- AとBから1つずつ球を取り出したとき、得点の和がk点になる取り出し方は何通りあるか

これの $2$ 点の場合を求めるには、以下のような計算をすることになる。

• $A$ から $0$ 点、$B$ から $2$ 点を取り出すのは、$5\times 1=5$ 通り
• $A$ から $1$ 点、$B$ から $1$ 点を取り出すのは、$3\times 2=6$ 通り
• $A$ から $2$ 点、$B$ から $0$ 点を取り出すのは、$2\times 6=12$ 通り
• これらを合計して、$5+6+12=23$ 通り

さて、この計算方法には見覚えがある。
これは数列の畳み込みそのものである。
つまり、母関数を利用して、以下の計算だけで $6$ つの答えが全て求まる。
$$
(5+3x+2x^2)\times (6+2x+x^2+x^3)=30+28x+23x^2+12x^3+5x^4+2x^5
$$

話を一般化しよう。
ある $2$ つの評価 $A$ と $B$ があり、それぞれ何点になるのが何パターンあるか、全て既知であるとする。
このとき、評価 $k$ になるパターンが $m$ 通りあることを、$mx^k$ と表現して母関数にする。
両者を掛け算して展開すれば、評価の和が何点になるのが何パターンあるか、全てわかるのである。

あるいは、評価が $3$ つ以上あってもよい。
評価が存在する個数分だけ母関数を用意して、全て掛け算すればよい。
例えば冒頭の問題は、$(x+x^2+x^3+x^4+x^5+x^6)^{20}$ を畳み込み $19$ 回で展開すれば解ける。

ただし、多数掛け算する場合は、計算量に注意。
畳み込みは長い数列ほど計算量がかかり、順に $1$ つずつ掛け算すると終盤の計算量が大変なことになる。
全体の中で短い方から $2$ つを畳み込むことを繰り返すのが効率が良い。
これは長さをコストとみなすハフマン符号を利用してハフマン木を構築することで実現できる。

が、さすがにこれだけのためにハフマン木を構築するのは手間がかかりすぎる。
そこで、ハフマン木よりは少し遅いが多くの場合は十分高速で、簡単に実装できる方法を紹介する。
数列をすべてqueueに入れる。
queueの先頭 $2$ つを取り出し、畳み込んで末尾に戻す。
これを、queueの長さが $1$ になるまで繰り返す。
$k$ 個の級数を掛ける場合、実質的にどのデータも $log k$ 回ずつしか扱われない。
長さ $2$ の級数 $k$ 個を掛けるのをFFT/NTTを用いて計算すると、計算量は $O(k \log^2k)$ となる。

指数型母関数の利用

通常型でない母関数を考えてみる。
すなわち、$a_n$ の値を、$x^n$ の係数ではなく、ある $d_n$ で割った $\dfrac{x^n}{d_n}$ の係数と思ってみる。
$$
A(x) = a_0\dfrac{x^0}{d_0} + a_1\dfrac{x^1}{d_1} + a_2\dfrac{x^2}{d_2} + a_3\dfrac{x^3}{d_3} + \dots + a_n\dfrac{x^n}{d_n} + \dots
$$

同じように $b_n$ の母関数 $B(x)$ も用意して掛け算してみる。
すると、畳み込みの結果の $\dfrac{x^n}{d_n}$ の係数は、次のようになる。
$$
c_n = \dfrac{d_n}{d_0d_n}a_0b_n + \dfrac{d_n}{d_1d_{n-1}}a_1b_{n-1} + \dfrac{d_n}{d_2d_{n-2}}a_2b_{n-2} + \dots + \dfrac{d_n}{d_kd_{n-k}}a_kb_{n-k} + \dots + \dfrac{d_n}{d_nd_0}a_nb_0
$$

つまり、畳み込む際に各数列の添字に応じた重みのようなものをつけることができる。
特に、$d_n=n!$ とした場合を指数型母関数と言い、畳み込み結果の係数が以下のようになる。
指数型というのは、全ての係数を $1$ としたときに、$A(x) = e^x$ と指数関数になるため。
$$
c_n = {}_n\mathrm{C}_0a_0b_n + {}_n\mathrm{C}_1a_1b_{n-1} + {}_n\mathrm{C}_2a_2b_{n-2} + \dots + {}_n\mathrm{C}_ka_kb_{n-k} + \dots + {}_n\mathrm{C}_na_nb_0
$$

つまり、$a_i$ と $b_j$ の積に、$i$ 個の $a$ と $j$ 個の $b$ を並び替える個数を掛け算する形になる。
これは $2$ つの積だけでなく、$3$ つ以上の積でも成り立つ。

この形は、以下のような問題で役に立つ。

7個のマスがあり、これらを赤、青、黄、緑の4色で塗り分ける。
ただし、同じ色は3回までしか使えない。
塗り方は何通りあるか。

この問題は、各色を使う個数を決めた後で、その並び順の個数を掛ける必要がある。
これが、指数型母関数の畳み込みのときにかかる数に都合よくなっているのである。
よって、$0$ 回から $3$ 回までの係数を $1$ として、指数型母関数は $A(x) = 1+x+\dfrac{1}{2}x^2+\dfrac{1}{6}x^3$ となる。
$4$ 色とも同じ母関数となるので、全体の母関数はその $4$ 乗で、以下となる。
$$
(1+x+\dfrac{1}{2}x^2+\dfrac{1}{6}x^3)^4 = 1+4x+8x^2+\dfrac{32}{3}x^3+\dfrac{21}{2}x^4+8x^5+\dfrac{29}{6}x^6+\dfrac{7}{3}x^7+O(x^8)
$$

この級数の $\dfrac{x^7}{7!}$ の係数を考えて、$\dfrac{7}{3}\times7!=11760$ が答え。
このように、選んだあとで並べ替えがある組み合わせ計算では、個数を指数にする指数型母関数が役立つ。

2変数の母関数の利用

文字を複数使い、$2$ 変数の級数にすることもできる。
例えば、毎ターンの行動を以下から選べるとする。

• 右へ $1$ マス移動する
• 右へ $2$ マス移動する
• 上へ $1$ マス移動する
• 右上へ $1$ マス移動する

このとき、右への移動を $x$、上への移動を $y$ に対応させると、母関数は以下のようになる。
$$
A(x,y) = x+x^2+y+xy
$$

では、このような移動を $5$ ターン行って、右に $5$ マス上に $2$ マス行った位置まで移動するパターン数は？
答えは、これの $5$ 乗における $x^5y^2$ の係数である $100$ 通り。

さて、数学的にはこれでいいとして、問題は畳み込みをどう行うか。
これは、十分大きな $k$ を用意して $y=x^k$ とおけばよい。
上の例では $5$ 乗した後での $x$ の指数は最大で $10$ である。
つまり $k=11$ として、$A(x) = x+x^2+x^{11}+x^{12}$ とすればよい。

切り落としをしたい場合は、切り落とす先を頑張って $0$ で埋めていくことになる。

分割統治による畳み込み

$1$ 方向にのみ進んでいく $1$ 次元の動的計画法で、一部のマスで値の加工をしたい場合がある。
このとき、仮に $8$ マスだとして、一般には以下のように配るDPで考えて処理する。

• 必要なら $0$ マス目の値を処理する
• $0$ マス目の値を、$1,2,3,4,5,6,7$ マス目に配る
• 必要なら $1$ マス目の値を処理する
• $1$ マス目の値を、$2,3,4,5,6,7$ マス目に配る
• 必要なら $2$ マス目の値を処理する
• $2$ マス目の値を、$3,4,5,6,7$ マス目に配る
• 必要なら $3$ マス目の値を処理する
• $3$ マス目の値を、$4,5,6,7$ マス目に配る
• 必要なら $4$ マス目の値を処理する
• $4$ マス目の値を、$5,6,7$ マス目に配る
• 必要なら $5$ マス目の値を処理する
• $5$ マス目の値を、$6,7$ マス目に配る
• 必要なら $6$ マス目の値を処理する
• $6$ マス目の値を、$7$ マス目に配る
• 必要なら $7$ マス目の値を処理する

この場合は、計算量が $O(N^2)$ かかる。

この処理順は、以下の $2$ つを守っている限り、並べ替えてもよい。

• $i$ マス目を処理する前に、それより前の全マスから $i$ マス目へ配り終わっている
• $i$ マス目を処理する前に、$i$ マス目からそれ後ろへのマスへ配ることはない

これを守る処理順として、以下のように組み替えることができる。

• 必要なら $0$ マス目の値を処理する
• $0$ マス目の値を、$1$ マス目に配る
• 必要なら $1$ マス目の値を処理する
• $0,1$ マス目の値を、$2,3$ マス目に配る
• 必要なら $2$ マス目の値を処理する
• $2$ マス目の値を、$3$ マス目に配る
• 必要なら $3$ マス目の値を処理する
• $0,1,2,3$ マス目の値を、$4,5,6,7$ マス目に配る
• 必要なら $4$ マス目の値を処理する
• $4$ マス目の値を、$5$ マス目に配る
• 必要なら $5$ マス目の値を処理する
• $4,5$ マス目の値を、$6,7$ マス目に配る
• 必要なら $6$ マス目の値を処理する
• $6$ マス目の値を、$7$ マス目に配る
• 必要なら $7$ マス目の値を処理する

これも、愚直にやれば計算量が $O(N^2)$ かかる。
しかし、この形だと、配る処理を畳み込みでまとめて処理できる。
具体的には、$2^k$ 個から $2^k$ 個に配る場合、$O(k2^k)$ で処理できる。
結果として、全体の計算量が、$O(N\log^2 N)$ になり、高速化される。

コード例は以下。
逆数のところで紹介した、AtCoder Libraryのconvolutionのラッパーを用いている。
また、DPテーブルの長さは $2$ の累乗になっている必要があることに注意。
f は、$k$ マス先への遷移のときに掛ける数を入れた配列。

// 0マス目の処理だけ先にする
for (int i=1; i<n; i++) {
  int w = i&-i;
  vector<long long> left = vector(dp.begin()+i-w,dp.begin()+i);
  vector<long long> transition = convolution_acl(left,f,2*w);
  for (int j=0; j<w; j++) {
    dp.at(i+j) += transition.at(w+j);
    dp.at(i+j) %= mod;
  }
  // iマス目の処理
}

注意点

特になし。

関連知識

順列組み合わせ

母関数では、組み合わせの数え上げを多項式の係数として表し、積や級数展開によって求める。

二項係数

分数型の母関数の級数展開や、母関数を用いた数え上げで現れる。

剰余類環

数論変換（NTT）では、素数を法とする剰余類環上で多項式の積を高速に求める。

繰り返し二乗法

数論変換（NTT）で必要となる累乗や逆元を高速に計算するために用いる。