Advertisement
matbensch

Untitled

Apr 30th, 2024 (edited)
652
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 2.03 KB | None | 0 0
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. // define pi
  5. const double PI = 3.14159265359;
  6.  
  7. // compute the DFT of x or the IDFT if invert
  8. void dft(vector<complex<double>>& x, bool invert = false) {
  9.     int N = x.size();
  10.     if (N <= 1) return;
  11.  
  12.     // spilt x into its even-indexed and odd-indexed terms
  13.     vector<complex<double>> E(N / 2), O(N / 2);
  14.     for (int i = 0; i < N / 2; i++) E[i] = x[2*i], O[i] = x[2*i+1];
  15.  
  16.     // compute the DFT/IDFT of each half    
  17.     dft(E, invert);
  18.     dft(O, invert);
  19.  
  20.     // w_k are the Nth roots of unity
  21.     complex<double> w(1);
  22.     double theta = 2 * PI / N * (invert ? -1 : 1);
  23.     complex<double> rot(cos(theta), sin(theta));
  24.    
  25.     // combine the halves
  26.     for (int k = 0; k < N / 2; k++) {
  27.         x[k] = E[k] + w * O[k];         // x[k] for k = 0...N/2-1
  28.         x[k + N/2] = E[k] - w * O[k];   // x[k+N/2] by periodicity
  29.         if (invert) x[k] /= 2, x[k + N / 2] /= 2;   // needed if inverting
  30.         w *= rot; // rotate w to the next root of unity
  31.     }
  32. }
  33.  
  34. // compute the coefficients of A(x)*B(x) where a and b are the coefficients of A(x) and B(x) respectively
  35. vector<int> poly_mul(vector<int>& a, vector<int>& b)
  36. {
  37.     // convert to complex data type
  38.     vector<complex<double>> ca(a.begin(), a.end());
  39.     vector<complex<double>> cb(b.begin(), b.end());
  40.     vector<int> res;
  41.  
  42.     // zero pad a and b until length is a power of 2
  43.     int N = 1;
  44.     while(N < ca.size() + cb.size()) N <<= 1;
  45.     ca.resize(N);
  46.     cb.resize(N);
  47.    
  48.     // compute the DFT of A and B
  49.     dft(ca);
  50.     dft(cb);
  51.  
  52.     // multiply the results and compute the IDFT
  53.     for(int i=0;i<N;i++) ca[i] *= cb[i];
  54.     dft(ca, true);
  55.    
  56.     // extract the coefficients and return
  57.     for(int i=0;i<N;i++) res.push_back(round(ca[i].real()));
  58.     while(res[res.size()-1]==0) res.pop_back();
  59.     return res;
  60. }
  61.  
  62. int main()
  63. {
  64.     vector<int> a = {-1, 2, 5};
  65.     vector<int> b = {1, 2};
  66.     vector<int> ab = poly_mul(a, b);
  67.    
  68.     for(int e: ab) cout << e << ' '; cout << '\n';
  69. }
  70.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement