Advertisement
nasarouf

FFT

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