fft.cpp

typedef long double T_DOUBLE; // May be changed to double for more speed and less precision

struct cpx
{
  cpx(){}
  cpx(T_DOUBLE aa):a(aa),b(0){}
  cpx(T_DOUBLE aa, T_DOUBLE bb):a(aa),b(bb){}
  T_DOUBLE a;
  T_DOUBLE b;
  T_DOUBLE modsq(void) const
  {
    return a * a + b * b;
  }
  cpx bar(void) const
  {
    return cpx(a, -b);
  }
};

cpx operator +(cpx a, cpx b)
{
  return cpx(a.a + b.a, a.b + b.b);
}

cpx operator *(cpx a, cpx b)
{
  return cpx(a.a * b.a - a.b * b.b, a.a * b.b + a.b * b.a);
}

cpx operator /(cpx a, cpx b)
{
  cpx r = a * b.bar();
  return cpx(r.a / b.modsq(), r.b / b.modsq());
}

cpx EXP(T_DOUBLE theta)
{
  return cpx(cos(theta),sin(theta));
}

const T_DOUBLE two_pi = 4 * acos((T_DOUBLE)0);

// in:     input array
// out:    output array
// step:   {SET TO 1} (used internally)
// size:   length of the input/output {MUST BE A POWER OF 2}
// dir:    either plus or minus one (direction of the FFT)
// RESULT: out[k] = \sum_{j=0}^{size - 1} in[j] * exp(dir * 2pi * i * j * k / size)
void FFT(cpx *in, cpx *out, int step, int size, int dir)
{
  if(size < 1) return;
  if(size == 1)
  {
    out[0] = in[0];
    return;
  }
  int half = size >> 1;
  FFT(in, out, step * 2, half, dir);
  FFT(in + step, out + half, step << 1, half, dir);
  for(int i = 0 ; i < half ; i++)
  {
    cpx even = out[i];
    cpx odd = out[i + half];
    out[i] = even + EXP(dir * two_pi * i / size) * odd;
    out[i + half] = even + EXP(dir * two_pi * (i + half) / size) * odd;
  }
}

// f:      input array f
// g:      input array g
// h:      output array h
// size:   length of f/g/h {MUST BE A POWER OF 2}
// RESULT: h[i] = sum of f[k]*g[i-k] for k=0..(size-1)
// Note: f and g are cyclic (e.g. f[-1] = f[size-1]), so they may be padded
// the right with zeroes to instead compute an acyclic convolution.
void FFTConvolution(cpx *f, cpx *g, cpx *h, int size)
{
  cpx *F = new cpx[size], *G = new cpx[size], *H = new cpx[size];
  FFT(f, F, 1, size, 1);
  FFT(g, G, 1, size, 1);
  for (int i = 0; i < size; i++) H[i] = F[i] * G[i];
  FFT(H, h, 1, size, -1);
  for (int i = 0; i < size; i++) h[i] = h[i] / size;
  delete [] F;
  delete [] G;
  delete [] H;
}