nasarouf

FFT

Mar 8th, 2017
228
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. //FFT
  2. //ref: AC Just a few more triangles! https://open.kattis.com/problems/moretriangles
  3.  
  4. #include <ccomplex>
  5. #include <complex>
  6. using namespace std;
  7.  
  8. typedef double ld;
  9. const ld PI = acos(-1.L);
  10. typedef complex<ld> Complex;
  11.  
  12. //precomputing these saves 50% time
  13. map<pair<int, int>, Complex> wns;
  14. void FFTinit(int N){
  15.     for (int i = 1; i <= N; i *= 2)
  16.         for (int dir = -1; dir <= 1; dir += 2)
  17.             wns[make_pair(i, dir)] = polar((ld)1, dir * 2 * PI / i);
  18. }
  19.  
  20. // *** Fast Fourier Transform (Recursive), ubc codearchive, heavily modified... ***
  21. //array pointers instead of vectors, calculations done in-place
  22. void rfft(int n, Complex* a, Complex* y, int dir = 1, int stride = 1) {
  23.     if (n == 1) { y[0] = a[0]; return; }
  24.     Complex wn = wns[make_pair(n, dir)], w = 1;
  25. //  Complex wn = polar((ld)1, dir * 2 * PI / n), w = 1; //precomputed instead
  26.     rfft(n / 2, a, y, dir, stride * 2);
  27.     rfft(n / 2, a + stride, y + n / 2, dir, stride * 2);
  28.     for (int k = 0; k < n / 2; k++, w *= wn) {
  29.         Complex y1 = y[k] + w*y[n / 2 + k];
  30.         Complex y2 = y[k] - w*y[n / 2 + k];
  31.         y[k + n / 2] = y2;
  32.         y[k] = y1;
  33.     }
  34. }
  35. //end code archive
  36. void ifft(int n, Complex* a, Complex* y) {
  37.     for (int i = 0; i < n; i++) a[i] /= ld(n);
  38.     rfft(n, a, y, -1);
  39. }
  40.  
  41. /** example:
  42.  
  43.     int N = 1 << int(ceil(log2(sz)));
  44.  
  45.     FFTinit(N);
  46.     vector<Complex> v(N, 0), c(N, 0), tmp(N, 0);
  47.  
  48.     //convolution
  49.     for (ll i = 0; i < n; i++) tmp[i] = v[i];
  50.     rfft(N, &tmp[0], &c[0]);
  51.     for (ll i = 0; i < N; i++) tmp[i] = c[i] * c[i];
  52.     ifft(N, &tmp[0], &c[0]);
  53.  
  54. */
Advertisement
Add Comment
Please, Sign In to add comment