【笔记】快速傅里叶变换（FFT）

OI-Wiki

#include 
#include 
#include 
//#include 
using namespace std;
const double Pi = acos(-1);
//typedef complex Comp;
const int MAXN = 3000006;
struct Comp {
	double x, y;
	Comp(double _x=0.0, double _y=0.0) {
		x = _x, y = _y;
	}
	Comp operator-(const Comp &b) const {
		return Comp(x-b.x, y-b.y);
	}
	Comp operator+(const Comp &b) const {
		return Comp(x+b.x, y+b.y);
	}
	Comp operator*(const Comp &b) const {
		return Comp(x*b.x-y*b.y, x*b.y+y*b.x);
	}
};
struct _FFT {
	int rev[MAXN];
	void change(Comp y[], int len) // len = 2^k
	{
		rev[0] = 0;
		for (int i=1; i< len; i++) {
			rev[i] = rev[i>>1]>>1;
			if (i&1) rev[i] |= len>>1;
		}
		for (int i=0; i< len; i++) if (i< rev[i]) swap(y[i], y[rev[i]]);
	}
	void fft(Comp y[], int len, int on) // on = 1 DFT, on = -1 IDFT
	{
		change(y, len);
		for (int h=2; h<=len; h<<=1) {
			Comp wn(cos(2*Pi/h), sin(on*2*Pi/h));
			for (int j=0; j< len; j+=h) {
				Comp w(1, 0); // w_n^k
				for (int k=j; k< j+(h>>1); k++) {
					Comp u = y[k], t = w*y[k+(h>>1)];
					y[k] = u + t, y[k+(h>>1)] = u - t;
					w = w * wn;
				}
			}
		}
		if (on==-1) for (int i=0; i< len; i++) y[i].x /= len;
	}
} FFT;
Comp F[MAXN], G[MAXN];
int main()
{
	int N, M; scanf("%d%d", &N, &M); N++, M++; // 0...N-1
	int len = 1; double a;
	while (len< (N<<1) || len< (M<<1)) len <<= 1;
	for (int i=0; i< N; i++) scanf("%lf", &a), F[i] = Comp(a, 0);
	for (int i=N; i< len; i++) F[i] = Comp(0, 0);
	for (int i=0; i< M; i++) scanf("%lf", &a), G[i] = Comp(a, 0);
	for (int i=M; i< len; i++) G[i] = Comp(0, 0);
	FFT.fft(F, len, 1), FFT.fft(G, len, 1);
	for (int i=0; i< len; i++) F[i] = F[i] * G[i];
	FFT.fft(F, len, -1);
	for (int i=0; i< N+M-1; i++) printf("%d ", (int)(F[i].x+0.5));
}