Senの競技プログラミング備忘録

こけた問題を自分用の解説で載せる。けんちょんさんのブログを目指したい。質的にも量的にも。こけた問題だけに限定するけど

高速フーリエ変換の超丁寧な解説! 実装のやり方や具体例も!

 前回の記事読了の方はお疲れ様でした。今回は高速フーリエ変換の実装とその具体例について書きます。


併せ読み *推奨* のAtcoder社の公式解説

実装

実装のやり方について書くと書きましたが、タイトル詐欺です。ごめんなさい。今回の記事の実装はAtcoderの解説の模擬コードをそのまま実装します。
ですので、模擬コードの中を見るだけでは理解できない人のための詳しい挙動を解説します。

実装する前に、思考の整理を

ひとまず、式の次数は 2n+1 = 2^a-1、2の累乗-1にします。
例えば、もともと5次式ならば、 0*x^6+0*x^7をつける感じで、とにかく空いてる次数は0で埋める感じでそろえます。理由はのちほど言います。

 f(x) = c_0 + c_1*x^1 + ... + c_{2n+1} * x^{2n+1} は、
 f_0(t) = c_0 + c_2 * t + .. + c_{2n}*t^n と  f_1(t) = c_1*t + c_3*t^2 + ... + c_{2n+1}*t^nを使って、
 f(x) = f_0(x^2)+x*f_1(x^2)でかけるんでした。
この時の f_0(t), f_1(t)の係数は偶数版ならば f_0(t)、奇数版なら f_1(t)に、低い次数から埋めていくのでした。

例えば、
 2x^2-x+3ならば
{3, -1, 2, 0} ->  3-1*x+2*x^2+0*x^3

{3, 2} {-1, 0} ->  3+2*x  -1+0*x

とこのように、子供に係数が伝播していきます。
そして、係数が1つになって伝播できない場合を考えます。
伝播途中の係数の列は左から多項式 x^i(0 -idx)にかかる係数でした。1つになったときは、つまり多項式ではなく、定数になるのです!
例えば、
{3, 2} ->  3*x^0+2*x^1
{3} {2} ->  3 * x^0  2 * x^0

二つの式に分割できるときは分けて、できないときはつまり、1つしか係数がないので定数となり、それを返せばよい。(変数自体が存在しない式なので、何を代入しようが定数になる)
あとは、再帰して計算した二つの半分に分けた列を f(x) = f_0(x^2)+x*f_1(x^2)で計算すれば、関数の返したい数列が求まります。
また、この構造をとると、再帰関数が完全二分木の各ノードのような形となるので、最初に足りない分を0埋めしてでも、次数を 2^n-1に合わせたのです。

実装例と解説

公式スライド25ページにあるものをC++に書き下ろしたものの解説をします。

typedef std::complex<double, double> Complex;
//実部も虚部もdoubleとなる複素数型を使う

const double PI = acos(-1);

vector<Complex> fft(vector<Complex> A, int n, int sgn = 1) {
        //係数の列Aを受け取って、1の(2^n)乗根の0乗から((2^n)-1)乗根までを、
        //列Aを係数とした多項式に代入した答えを返す
        //1のN乗根は複素数になるので、係数は複素数になる たとえ最終的な答えが実数であったとしても

	if (n == 1)return A;
        //前での説明通り、1つしか係数がない式は定数なので、それを返す
	
        vector<Complex> f0, f1;
	for (int i = 0; i < (n / 2); i++) {
		f0.push_back(A[i * 2 + 0]);
		f1.push_back(A[i * 2 + 1]);
	}
        //奇数番目と偶数番目に係数を分けて、二つのサイズが半分の係数列に分ける。

	f0 = fft(f0, n / 2, sgn), f1 = fft(f1, n / 2, sgn);
        //解説通り、二つの係数列となるものを再帰的に求める。DFSの順番で(再帰関数を使ってるのでそれはそう)
        //もともとf0, f1は係数列だけど、それを関数の答え、つまり二つの半分のサイズの列に代入した答えの配列に再利用

	Complex zeta = Complex(cos(2.0 * PI / n), sin(2.0 * PI / n) * sgn);
	Complex pow_zeta = 1;
        //(2^n)-1次の式となるので、(2^n)個の1のn乗根を用意して計算しなければならない。

	for (int i = 0; i < n; i++) {
		A[i] = f0[i % (n / 2)] + pow_zeta * f1[i % (n / 2)];
		pow_zeta *= zeta;
        }
        //あらかじめ再帰して求めた二つのサイズ半分の係数列の関数に、1のn乗根を代入して、f(x)=f_0(x^2)+x*f_1(x)のとおりに計算する。
        //再帰で求めた配列は2^(n - 1)の長さで、Atcoder社のスライドの考察通り、後半部分は前半部分と同じなので、このように計算してる。
	return A;
}

sgnというものがありますが、今はひとまずsgn = 1であるという認識で構いません。逆フーリエ変換の時に説明します。

経過の具体例

f:id:Sen_comp:20191002234311p:plain

大字の漢数字は順番を表してます。係数列でもともと値を持ったのは赤代入して計算した結果の数列は緑で表示した。

訂正! 引数が{-3, 2, 1, 0}の再帰関数のreturn の部分は{0, -4-2i, 4, -4+2i}となっていますが、
正しくは{0, -4-2i, -4, -4+2i}です!(3項目にマイナスをつけ忘れました)

また、変換前の関数を f(x)、変換後を \hat{f}(x)、変換後計算に使う偶数項から成る多項式 f_0(x)、奇数項は f_1(x)とします。これらは、同じ再帰関数内の呼称です!(再帰関数Aの中での f_0(x)再帰関数Bの中での f(x)になりえます。)

  • (壱) この再帰では f(x)=x^2+2x-3を変換します。
  • (弐) 偶数次の項の係数{-3, 1}と奇数次の項の係数{2, 0}に分けます。それぞれ f_0(x)=x-3 f_1(x)=2となります。
  • (参) 引数が{-3, 2, 1, 0}の関数からの再帰関数の呼び出し。呼び出し元の偶数部に相当し、この関数内での f(x) = x-3がとなる。
  • (肆) 偶数次の項の係数{-3}と奇数次の項の係数{-1}に分けます。それぞれ f_0(x)=-3 f_1(x)=2となります。
  • (伍) 係数が1つしかないので、 f(x)=-3の定数関数となる。よって、何が xに代入されても一定なので、-3を返します。
  • (陸) 係数が1つしかないので、 f(x)=1の定数関数となる。よって、何が xに代入されても一定なので、1を返します。
  • (質) 再帰から得た二つのハーフサイズの数列の計算をします。

 f(\zeta_2^0) = -3 + \zeta_2^0*1=-2
 f(\zeta_2^1) = -3 + \zeta_2^1*1=-4

  • (捌) 計算済みの数列を \hat{f}(x)に1のn乗根を2^n個代入て求めた値として、返します。
  • (玖) 引数が{-3, 2, 1, 0}の関数からの再帰関数の呼び出し。呼び出し元の奇数部に相当し、この関数内での f(x) = 2がとなる。
  • (拾) 偶数次の項の係数{2}と奇数次の項の係数{0}に分けます。それぞれ f_0(x)=2 f_1(x)=0となります。
  • (拾壱) 係数が1つしかないので、 f(x)=2の定数関数となる。よって、何が xに代入されても一定なので、2を返します。
  • (拾弐) 係数が1つしかないので、 f(x)=0の定数関数となる。よって、何が xに代入されても一定なので、0を返します。
  • (拾参) 再帰から得た二つのハーフサイズの数列の計算をします。

 f(\zeta_2^0) = 2 + \zeta_2^0*0=2
 f(\zeta_2^1) = 2 + \zeta_2^1*0=2

  • (拾肆) 計算済みの数列を \hat{f}(x)に1のn乗根を2^n個代入て求めた値として、返します。
  • (拾伍) 再帰から得た二つのハーフサイズの数列の計算をします。

 f(\zeta_4^0) = -2 + \zeta_4^0*2=0
 f(\zeta_4^1) = -4 + \zeta_4^1*2=-4-2i
 f(\zeta_4^1) = -2 + \zeta_4^2*2=-4
 f(\zeta_4^1) = -4 + \zeta_4^3*2=-4+2i

  • (拾陸) 計算済みの数列を \hat{f}(x)に1のn乗根を2^n個代入て求めた値として、返します。

これで、 f(x)=x^2+2x-3に相当する \hat{f}(x)=(-4+2i)*x^3+(-4)*x^2+(-4-2i)*x+0が求まりました!

フーリエ変換とsgnの意味

先ほど説明の中で放置していたsgnについて説明します。
フーリエ変換は、 f(x)に対して、1のn乗根をすべて順番に代入し、その関数の計算結果を係数とする \hat{f}(x)を作る変換でした。

フーリエ変換はその逆で、 \hat{f}(x)に対して、1のn乗根の逆数をすべて順番に代入し、その関数の計算結果を係数とする f(x)を作る変換です。フーリエ変換、逆フーリエ変換を行えば、Atcoderのスライドのとおり、 f(x)の係数を抽出できます。

ところで、フーリエ変換は、代入する数が1のn乗根が1のn乗根の逆数になった以外はフーリエ変換と同じです!
なので、先ほど描いたフーリエ変換の関数を流用できます。sgnが-1の時は、1のn乗根の逆数が代入されるようになっています。

ソースコードの全体的な流れ

Atcoder社の解説の流れのとおりに書きあげました。実質的にはほぼほぼtsutajさんのソースコードのパクリです。

const double PI = acos(-1);

vector<Complex> fft(vector<Complex> A, int n, int sgn = 1) {
	//中略
}

vector<Complex> inv_fft(vector<Complex> A, int n) {
	auto ret = fft(A, n, -1);
	for (int i = 0; i < n; i++) {
		ret[i] /= n;
	}
	return ret;
}

vector<Complex> multiply(vector<Complex> &X, vector<Complex> &Y) {
	int n = 1;
	while (n < (X.size() + Y.size() + 1)) n *= 2;

	vector<Complex> ret;

	X.resize(n), Y.resize(n);
	X = fft(X, n, 1),
	Y = fft(Y, n, 1);

	for (int i = 0; i < n; i++) {
		ret.push_back(X[i] * Y[i]);
	}
	return inv_fft(ret, n);
}

最後に

ここまで読んでいただきありがとうございました。私も2記事、合計11300字前後を書いて、フーリエ変換への理解が非常に進みました。この記事が所学者の一助になれば幸いです。

最後に、Atcoder社の解説さん(解説スライドの追加解説の許可をくれたchokudaiさん)、ソースコードを始めに色々参考にさせてもらったtsutajさん、鈍い頭を持つ僕にわからせようとしたhotman78さんに大変感謝します。

以上。