FFT（NTT） 学习笔记

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/qq_33468963/article/details/78587792

inline void fft(comp a[],int f) {
	for (i=0;i<N;i++) if (i < rev[i]) swap(a[i],a[rev[i]]);
	for (i=1;i<N;i<<=1) {
		comp wn(cos(M_PI/i),f*sin(M_PI/i)),w,X,Y;
		for (j=0;j<N;j+=(i<<1)) {
			w = comp(1,0);
			for (k=0;k<i;k++,w = w * wn) {
				X = a[j+k]; Y = w*a[j+i+k];
				a[j+k] = X+Y;  a[j+i+k] = X-Y; 
			}
		}
	}
	if (f == -1) for (i=0;i<N;i++) a[i].r /= N;
}

typedef vector<int> poly;
inline void dft(int n,comp *a,int f) {
 for (int i=1;i<n;i++) {
  rev[i] = (rev[i>>1] >> 1) | ((i&1) ? (n>>1) : 0);
  if (i > rev[i]) swap(a[i],a[rev[i]]);
 }
 for (int i=1;i<n;i<<=1) {
  comp w, wn((long double)cos(PI/i),(long double)f*sin(PI/i)), X, Y;
  for (int j=0;j<n;j+=(i<<1)) {
   w = comp(1,0);
   for (int k=0;k<i;k++,w = w * wn) {
    X = a[j+k], Y = a[j+k+i] * w;
    a[j+k] = X+Y;  a[j+k+i] = X-Y;
   }
  }
 }
 if (f == -1) for (int i=0;i<n;i++) a[i].re /= n;
}
inline poly mtt(poly x,poly y) {
 int n;
 for (n=1;n <= x.size()*2 && n <= y.size()*2;n<<=1);
 for (int i=0;i<n;i++) k1[i] = k2[i] = b1[i] = b2[i] = comp(0,0);
 
 for (int i=0;i<x.size();i++) k1[i].re = x[i] >> 15, b1[i].re = x[i] & 0x7fff;
 for (int i=0;i<y.size();i++) k2[i].re = y[i] >> 15, b2[i].re = y[i] & 0x7fff;
 dft(n,k1,1);  dft(n,k2,1);  dft(n,b1,1);  dft(n,b2,1);
 for (int i=0;i<n;i++) {
  a[i] = k1[i] * k2[i];
  b[i] = k1[i] * b2[i] + k2[i] * b1[i];
  c[i] = b1[i] * b2[i];
 }
 dft(n,a,-1);  dft(n,b,-1);  dft(n,c,-1);
 
 poly z(x.size() + y.size() - 1);
 for (int i=0;i<z.size();i++) {
  z[i] = ((fl(a[i].re) % P) * ((1ll<<30) % P) + (fl(b[i].re) % P) * (1ll<<15) + (fl(c[i].re) % P)) % P;
 }
 return z;
}