peltorator

FFT

Jul 15th, 2017
160
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. typedef long double ld;
  2. typedef complex<ld> cmpl;
  3.  
  4. const ld PI = atan2(0, -1);
  5.  
  6. int rev(int i, int k)
  7. {
  8.     int j = 0;
  9.     for (int l = 0; l < k; l++)
  10.         j += (((i >> l) & 1) << (k - l - 1));
  11.     return j;
  12. }
  13.  
  14. void fft(vector<cmpl> &a, bool invert)
  15. {
  16.     int n = (int)a.size();
  17.     int k = 0;
  18.     while ((1 << k) < n)
  19.         k++;
  20.     for (int i = 0; i < n; i++)
  21.     {
  22.         int j = rev(i, k);
  23.         if (j < i)
  24.             swap(a[i], a[j]);
  25.     }
  26.     for (int l = 1; l < n; l *= 2)
  27.     {
  28.         ld al = (ld)(invert ? -1 : 1) * PI / (ld)(l);
  29.         cmpl wn = cmpl(cos(al), sin(al));
  30.         for (int i = 0; i < n; i += 2 * l)
  31.         {
  32.             cmpl w = cmpl(1, 0);
  33.             for (int j = 0; j < l; j++)
  34.             {
  35.                 cmpl y = a[i + j], z = a[i + j + l] * w;
  36.                 a[i + j] = y + z;
  37.                 a[i + j + l] = y - z;
  38.                 w *= wn;
  39.             }
  40.         }
  41.     }
  42.     if (invert)
  43.     {
  44.         for (int i = 0; i < n; i++)
  45.             a[i] /= (ld)n;
  46.     }
  47. }
  48.  
  49. const ld eps = 1e-2;
  50.  
  51. ld Abs(ld x)
  52. {
  53.     if (x < 0)
  54.         return -x;
  55.     return x;
  56. }
  57.  
  58. vector<ll> mult(vector<ll> a, vector<ll> b)
  59. {
  60.     if (!a.size() || !b.size())
  61.         return vector<ll>(1, 0);
  62.     int k = 1;
  63.     while (k <= a.size() + b.size() + 5)
  64.         k *= 2;
  65.     vector<cmpl> x(k, 0), y(k, 0);
  66.     for (int i = 0; i < a.size(); i++)
  67.         x[i] = a[i];
  68.     for (int i = 0; i < b.size(); i++)
  69.         y[i] = b[i];
  70.     fft(x, 0);
  71.     fft(y, 0);
  72.     for (int i = 0; i < k; i++)
  73.         x[i] *= y[i];
  74.     fft(x, 1);
  75.     vector<ll> c(k);
  76.     for (int i = 0; i < k; i++)
  77.     {
  78.         ld cur = x[i].real();
  79.         ll bst = cur - 5;
  80.         for (ll ch = cur - 5; ch < cur + 5; ch++)
  81.             if (Abs((ld)bst - cur) > Abs((ld)ch - cur))
  82.                 bst = ch;
  83.         c[i] = bst;
  84.     }
  85.     while (c.size() > (int)a.size() + (int)b.size() - 1)
  86.         c.pop_back();
  87.     return c;
  88. }
RAW Paste Data