【学习小结】FFT、NTT模板和原理

FFT、NTT原理

from zhouzhendong

很好的总结

from coco_T

模板

namespace fft { 
int base = 1;
poly rev = {0, 1};
vector<Complex> roots = {Complex(1,0),Complex(1,0)};

void ensure_base(int nbase) { //所有dft需要的预处理
  if (nbase <= base) {
    return;
  }
  rev.resize(1 << nbase);
  for (int i = 0; i < 1 << nbase; ++i) { //预处理翻转位
    rev[i] = rev[i >> 1] >> 1 | (i & 1) << (nbase - 1);
  }
  roots.resize(1 << nbase);
  while (base < nbase) { //预处理单位根
    Complex z(cos(Pi / (1 << base)),sin(Pi / (1 << base)));
    for (int i = 1 << (base - 1); i < 1 << base; ++i) { //处理1 << (base + 1)次单位根
      roots[i << 1] = roots[i];
      roots[i << 1 | 1] = roots[i] * z;
    }
    ++base;
  }
}

void dft(vector<Complex> &a) {
  int n = a.size(), zeros = __builtin_ctz(n);
  ensure_base(zeros);
  int shift = base - zeros;
  for (int i = 0; i < n; ++i) {
    if (i < rev[i] >> shift) {
      swap(a[i], a[rev[i] >> shift]);
    }
  }
  for (int i = 1; i < n; i <<= 1) {
    for (int j = 0; j < n; j += i << 1) {
      for (int k = 0; k < i; ++k) {
        Complex x = a[j + k], y = a[j + k + i] * roots[i + k];
        a[j + k] = x + y;
        a[j + k + i] = x - y;
      }
    }
  }
}
vector <Complex> tmpa,tmpb;
poly multiply(poly a, poly b) { //先转化成complex再运算，注意取模，注意resize的问题
  int need = a.size() + b.size() - 1, nbase = 0;
  while (1 << nbase < need) {
    ++nbase;
  }
  ensure_base(nbase);
  int sz = 1 << nbase;
  tmpa.resize(sz) , tmpb.resize(sz);
  fill(tmpa.begin(),tmpa.end(),Complex(0,0)) , fill(tmpb.begin(),tmpb.end(),Complex(0,0));
  rep(i,0,a.size() - 1) tmpa[i] = Complex(a[i],0);
  rep(i,0,b.size() - 1) tmpb[i] = Complex(b[i],0);

  dft(tmpa) , dft(tmpb);

  rep(i,0,sz - 1) tmpa[i] *= tmpb[i];

  dft(tmpa);
  reverse(tmpa.begin() + 1,tmpa.end());
  a.resize(need);
  rep(i,0,need - 1) a[i] = (ll)(tmpa[i].real() / sz + 0.5) % mod;
  return a;
}

}

using fft::multiply;

poly& operator *= (poly &a, const poly &b) {
  if ((int) min(a.size(), b.size()) < 128) {
    poly c = a;
    a.assign(a.size() + b.size() - 1, 0);
    for (int i = 0; i < (int) c.size(); ++i) {
      for (int j = 0; j < (int) b.size(); ++j) {
        add(a[i + j], mul(c[i], b[j]));
      }
    }
  } else {
    a = multiply(a, b);
  }
  return a;
}

poly operator * (const poly &a, const poly &b) {
  poly c = a;
  return c *= b;
}