jinya-blog: mod p での convolution を FFT で実装する（備忘録）

mod p でのたたみ込みを FFT で実装するコードを自作するための備忘録。

フーリエ変換のまとめ

離散フーリエ変換

$h$ をある体$F$の上での$x$の関数

\begin{equation*}h(x) = \sum_{i=0}^{n-1} c_i x^i\end{equation*}

とする。また、$F$には1の原始$n$乗根が存在すると仮定し、それを$\zeta_n$とする。つまり$\zeta_n$は

\begin{equation*} \zeta_n^i \left\{ \begin{array}{} = 1 & \text{ if } i = n \\ \ne 1 & \text{ if } 1 \le i \le n-1 \end{array} \right.\end{equation*}

なる$F$の元とする。

すると、$h(\zeta_n^i)$はまた$F$の元なので、これを$\hat{c}_i$と置き（つまり $\hat{c}_i = h(\zeta_n^i)$）、関数$\hat{h}$を

\begin{equation*}\hat{h}(x) = \sum_{i=0}^{n-1} \hat{c}_i x^i\end{equation*}

とすると、$h \mapsto \hat{h}$ は$F[x]$から$F[x]$への写像となる。

関数の積と離散フーリエ変換

さて、ここで多項式環 $F[x]$ において、$f(x)$ と $g(x)$ との積を考える。$f, g$ をそれぞれ

\begin{equation*} f(x) = \sum_{i=0}^{n-1}a_i x^i , \quad g(x) = \sum_{j=0}^{n-1}b_j x^j\end{equation*}

と表すが、簡単のため、$n$は十分大きい数とし、$a_i$, $b_j$ は$i, j$が大きいと$0$であるようにとっておく。特に、$a_i$が$0$でない最大の$i$と、$b_j$が$0$でない最大の$j$との和が$n$未満となるように、$n$を十分大きくしておく。すると、

\begin{equation*} f(x) . g(x) = \left( \sum_{i=0}^{n-1}a_i x^i \right)\left( \sum_{i=0}^{n-1}b_i x^i \right) = \sum_{k=0}^{n-1} \left( \sum_{i+j=k} a_i b_j \right) x^k \end{equation*}

となる。そこでこの右辺の関数を $f$と$g$との積とし、$f*g$と書くことにする。

では、この関数の積は先の写像$h \mapsto \hat{h}$でどうなるか。

\begin{equation*} (f*g)(x) \mapsto \widehat{(f*g)}(x) = \sum_{k=0}^{n-1} (f*g)(\zeta_n^k) x^k \end{equation*}

だが、$(f*g)(x)$の定義は$(f*g)(x)=f(x).g(x)$だったので、$(f*g)(\zeta_n^k)=f(\zeta_n^k).g(\zeta_n^k)$、よって、

\begin{equation*} (f*g)(x) \mapsto \widehat{(f*g)}(x) = \sum_{k=0}^{n-1} f(\zeta_n^k)g(\zeta_n^k) x^k = \sum_{k=0}^{n-1} \hat{a}_k \hat{b}_k x^k \end{equation*}

となる。つまり、$F[x]$での関数の積は、$h \mapsto \hat{h}$で移した先では同じ次数の係数の積に変換されていた。

離散フーリエ逆変換

もう一つよく似た写像で、

\begin{equation*} h(x) \mapsto \check{h}(x) = \sum_{i=0}^{n-1} h(\zeta_n^{-i}) x^i \end{equation*}

を作る。この写像は、$\zeta_n^i$ が$\zeta_n^{-i}$に置き換わったものであるが、$\hat{h}$をこの写像で飛ばすと

\begin{equation*} \hat{h}(x) \mapsto \check{\hat{h}}(x) = \sum_{i=0}^{n-1} \hat{h}(\zeta_n^{-i}) x^i \end{equation*}

ここで、

\begin{equation*} \hat{h}(\zeta_n^{-i}) = \sum_{j=0}^{n-1} h(\zeta_n^j) (\zeta_n^{-i})^j = \sum_{j=0}^{n-1} \sum_{k=0}^{n-1} c_k (\zeta_n^j)^k (\zeta_n^{-i})^j = \sum_{k=0}^{n-1} c_k \sum_{j=0}^{n-1} \zeta_n^{j(k-i)} \end{equation*}

すると、最後の$\sum_{j=0}^{n-1}(\zeta_n^{k-i})^j$ は、$k=i$なら$n$、$k\ne i$なら$0$となる。なぜなら、まず$k=i$ならば$\zeta_n^{k-i}=1$で自明なので、$k\ne i$とすると、$-(n-1) \le k - i \le n - 1$であり、$\zeta_n$は1の原始$n$乗根なので、これらはすべて$1$でない$1$の$n$乗根となっている。よって、$\zeta_n^{k-i}=\xi$とし、$\sum_{j=0}^{n-1}\xi^j$ に $\xi - 1$をかけると、

\begin{equation*} (\xi - 1)\sum_{j=0}^{n-1}\xi^j = (\xi - 1)(\xi^{n-1} + \xi^{n-2} + ...+ \xi + 1) = \xi^n - 1 \end{equation*}

$\xi$は1の$n$乗根なので右辺は0、また、左辺の$(\xi -1)$は0ではないので、$\sum_{j=0}^{n-1}\xi^j = 0$であることがわかった。

以上により、

\begin{equation*} \hat{h}(\zeta_n^{-i}) = n c_i , \quad \check{\hat{h}}(x) = n \sum_{i=0}^{n-1} c_i x^i = n h(x)\end{equation*}

を得た。つまり、$h \mapsto \hat{h} \mapsto \check{\hat{h}} = n h$であり、特にこの二つの写像は全単射であることもわかった。

FFT（高速フーリエ変換）

さて、関数$h(x)=\sum_{k=0}^{n-1}c_k x^k$が与えられたとき、この離散フーリエ変換を求めたいとする。つまり、$h(\zeta_n^i)$をすべての$i$について求めたいのだが、計算式通りに求めようとすると、$0 \le i \le n-1$ と $0 \le k \le n-1$との組み合わせ分、$n\times n$回の計算が必要になる。つまり、計算量は$\mathcal{O}(n^2)$であるので、$n > 10^5$などになってくるとかなり厳しい。

そこで、この計算を$\mathcal{O}(n \log n)$程度で済ませるのがFFT。

まず、$n$を適当に大きくとって$2$のべき乗になるようにし、増えた分の$c_k$は0で埋める。（2のべき乗でなくても方法はあるが、2のべき乗が最も高速に処理できる。）

次に、$h$の各項を偶数番目と奇数番目に分けて、

\begin{equation*} h(x) = \sum_{k=0}^{n/2-1}c_{2k}x^{2k} + x \sum_{k=0}^{n/2-1}c_{2k+1}x^{2k} \end{equation*}

とし、それぞれの項を$h_0$、$h_1$とする。つまり、

\begin{eqnarray*} h_0(x) &=& \sum_{k=0}^{n/2-1}c_{2k}x^k \\ h_1(x) &=& \sum_{k=0}^{n/2-1}c_{2k+1}x^k \end{eqnarray*}

とおけば、

\begin{eqnarray*}h(x) = h_0(x^2) + x h_1(x^2)\end{eqnarray*}

なので、特に

\begin{eqnarray*}h(\zeta_n^i) = h_0(\zeta_{n/2}^i) + \zeta_n^i h_1(\zeta_{n/2}^i)\end{eqnarray*}

である。$h_0$, $h_1$の中に入っているのは $(\zeta_n^i)^2 = \zeta_{n/2}^i$、つまり、1の原始$n/2$乗根 $\zeta_n^2$ の$i$乗であることに注意する。

さて、この$h_0$や$h_1$と、$0 \le i \le n-1$について、$h_0(\zeta_{n/2}^i)$および$h_1(\zeta_{n/2}^i)$を求めれば、$h(\zeta_n^i)$がわかるのだが、$\zeta_{n/2}^i$は1の原始$n/2$乗根のべき乗なので、$0 \le i \le n/2 -1$まで計算すれば、あとはそれが再度繰り返されるだけ。つまり、$h_0$および$h_1$からそれぞれ$h_0(\zeta_{n/2}^i)$, $h_1(\zeta_{n/2}^i)$を求める作業は$\mathcal{O}(2\times(n/2)^2)=\mathcal{O}(n^2/2)$かかり、そこから$h$再構成するのに$\mathcal{O}(n)$かかるので、合わせて$\mathcal{O}(n + n^2/2)$かかる。

そして、$h$を$h_0$と$h_1$とに分割した作業をさらに$h_0$と$h_1$にも適用すれば、計算時間は $\mathcal{O}(n + 2 \times n/2 + n^2/4)$に減る。これを繰り返し、2のべき乗であった$n$が1になるまで分割すれば、結局計算量は$\mathcal{O}(n\log n )$程度になる。

Convolutionとの関係

数列$\{a_i\}_{i=0}^{n-1}, \{b_i\}_{i=0}^{n-1}$が与えられたとき、これを係数に持つ関数$f(x), g(x)$をつくって積を取れば、

\begin{equation*} f(x) . g(x) = \left( \sum_{i=0}^{n-1}a_i x^i \right)\left( \sum_{i=0}^{n-1}b_i x^i \right) = \sum_{k=0}^{n-1} \left( \sum_{i+j=k} a_i b_j \right) x^k \end{equation*}

となるが、この式の右辺の各$x^k$の係数に、数列$\{a_i\}$ と $\{b_j \}$ とのすべての$i+j=k$についての convolution が出現する。つまり、$h = f * g$とすると、$c_k = \sum_{i+j=k}a_i b_j$ となるので、Convolutionの目的は$h$の係数$c_k$を全ての$k$について求めることにある。

ところで、$h = f*g$ を離散フーリエ変換で飛ばすと、$\hat{h}$の係数は

\begin{equation*} h(\zeta_n^i) = f(\zeta_n^i) \times g(\zeta_n^i) \end{equation*}

である。つまり、離散フーリエ変換で飛ばした先では、関数の積は係数の積に置き換わるので、$\hat{h}$の係数は$\hat{f}$ および $\hat{g}$ の係数から$\mathcal{O}(n)$で求まる。

そこで、本来求めたかった$h$の係数$c_i$を、一旦フーリエ変換した先の係数の積で$\hat{h}$を求めておいてから、これをフーリエ逆変換で戻してやれば（逆変換も$\mathcal{O}(n \log n)$で完了するので）、$c_i$の全てを$\mathcal{O}(n\log n)$で求めることができる。

\begin{array}{cccc} f,g &\xrightarrow{\mathcal{O}(n \log n)} & \hat{f},\hat{g} & \\ \downarrow & &\downarrow & \mathcal{O}(n) \\ \check{\hat{h}} = nh &\xleftarrow{\mathcal{O}(n \log n)} & \hat{h} & \end{array}

mod p での Convolution

以上の操作を、素数$p$に対する体$F=\mathbb{Z}/p\mathbb{Z}$で行う。このときポイントになるのは、なるべく大きい2のべき乗$n$について、1の原始$n$乗根が存在することである。逆に言えば、そのような$n$が存在しない場合には、（簡単な）FFTを使って$\mathbb{Z}/p\mathbb{Z}$上のconvolutionを高速に行うことは難しくなる。では、このような$n$が取れる条件は何か。

素数$p$に対する$\mathbb{Z}/p\mathbb{Z}$の乗法群$(\mathbb{Z}/p\mathbb{Z})^\times$は、位数が$p-1$の群になる。よって、すべての$a \in \mathbb{Z}/p\mathbb{Z}$ について、$a \ne 0$ならば $a^{p-1} = 1 \mod p$。特に、$(\mathbb{Z}/p\mathbb{Z})^\times$は位数が$p-1$の巡回群にもなるので、必ず位数が$p-1$の元が存在する。それを$r$とすると、$1 \le i < p-1$について$r^i \ne 1 \mod p$ かつ $g^{p-1} = 1 \mod p$である。また、0でない全ての元の位数は$p-1$の約数になっていることもわかるので、特に、1の原始根があるとすれば、その位数は$p-1$の約数である。

さて、競プロで素数の例としてよく使われるのは、$10^9+7$や$998244353$であるが、それぞれについて乗法群の位数を素因数分解すると、前者は$10^9+7 - 1 = 2\times (5 \times 10^8 + 3)$、後者は $998244353 - 1 = 2^{23}\times 7\times 17$である。つまり、前者は乗法群の位数がまた巨大な素数を素因数としてもっており、他方後者は非常に大きい2のべき乗をもっている。上で述べたように、FFTは大きな2のべき乗を$n$として用いたいので、$p = 10^9+7$ では難しく、$p = 998244353$はふさわしいことがわかる。実際、高難易度の大会でない限りは、convolutionの問題では後者の$p$が用いられることがほとんどであろう。

以下、$p=998244353$で考える。次はこの$p$を用いた$\mathbb{Z}/p\mathbb{Z}$にて1の原始根を探す必要があるのだが、実は$3$は1の原始$p-1$乗根の一つであることがわかる。これは、$3^i \mod p$を順に全部計算してみればよい。$3$の位数は$p-1 = 2^{23}\times 7 \times 17 = 2^{23}\times 119$なので、$3^{119} \mod p = 15311432$は1の原始$2^{23}$乗根の一つである。

あとは、$15311432$を2乗していくことで、1の原始$2^{22}$乗根$267099868$、1の原始$2^{21}$乗根$733596141$、・・・を得られる。

以上でFFTを使ったConvolutionを実装する手がかりが全て得られた。