matbensch

CP-Algorithms FFT Implementation

Feb 2nd, 2023
248
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 1.19 KB | None | 0 0
  1. // from https://cp-algorithms.com/algebra/fft.html#implementation
  2. using cd = complex<double>;
  3. const double PI = acos(-1);
  4.  
  5. void fft(vector<cd> & a, bool invert) {
  6.     int n = a.size();
  7.     if (n == 1)
  8.         return;
  9.  
  10.     vector<cd> a0(n / 2), a1(n / 2);
  11.     for (int i = 0; 2 * i < n; i++) {
  12.         a0[i] = a[2*i];
  13.         a1[i] = a[2*i+1];
  14.     }
  15.     fft(a0, invert);
  16.     fft(a1, invert);
  17.  
  18.     double ang = 2 * PI / n * (invert ? -1 : 1);
  19.     cd w(1), wn(cos(ang), sin(ang));
  20.     for (int i = 0; 2 * i < n; i++) {
  21.         a[i] = a0[i] + w * a1[i];
  22.         a[i + n/2] = a0[i] - w * a1[i];
  23.         if (invert) {
  24.             a[i] /= 2;
  25.             a[i + n/2] /= 2;
  26.         }
  27.         w *= wn;
  28.     }
  29. }
  30.  
  31. vector<int> multiply(vector<int> const& a, vector<int> const& b) {
  32.     vector<cd> fa(a.begin(), a.end()), fb(b.begin(), b.end());
  33.     int n = 1;
  34.     while (n < a.size() + b.size())
  35.         n <<= 1;
  36.     fa.resize(n);
  37.     fb.resize(n);
  38.  
  39.     fft(fa, false);
  40.     fft(fb, false);
  41.     for (int i = 0; i < n; i++)
  42.         fa[i] *= fb[i];
  43.     fft(fa, true);
  44.  
  45.     vector<int> result(n);
  46.     for (int i = 0; i < n; i++)
  47.         result[i] = round(fa[i].real());
  48.     return result;
  49. }
  50.  
Advertisement
Add Comment
Please, Sign In to add comment