PloadyFree

Polynomial multiplication with FFT

Feb 21st, 2019
196
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Java 3.77 KB | None | 0 0
  1.     class Polynomial {
  2.         long[] coefficients;
  3.  
  4.         Polynomial(long[] coefficients) {
  5.             this.coefficients = coefficients;
  6.         }
  7.  
  8.         Polynomial multiply(Polynomial other) {
  9.             long[] coef = new FFT().multiply(coefficients, other.coefficients);
  10.             return new Polynomial(Arrays.copyOfRange(coef, 0, coefficients.length * 2 - 1));
  11.         }
  12.  
  13.         @Override
  14.         public String toString() {
  15.             StringJoiner sj = new StringJoiner(" ");
  16.             for (long x : coefficients) sj.add("" + x);
  17.             return sj.toString();
  18.         }
  19.     }
  20.  
  21.     class FFT {
  22.         void fft(double[] real, double[] imag, boolean inv) {
  23.             int n = real.length;
  24.             int m = 0;
  25.             while (1 << m < n) m++;
  26.             int[] rev = new int[n];
  27.             for (int i = 1, j = -1; i < n; i++) {
  28.                 if ((i & (i - 1)) == 0) j++;
  29.                 rev[i] = rev[i ^ (1 << j)] | (1 << (m - 1 - j));
  30.                 if (i < rev[i]) {
  31.                     swap(real, i, rev[i]);
  32.                     swap(imag, i, rev[i]);
  33.                 }
  34.             }
  35.  
  36.             for (int len = 1; len < n; len <<= 1) {
  37.                 double root = Math.PI / len;
  38.                 if (inv) root = -root;
  39.                 final double rootReal = Math.cos(root);
  40.                 final double rootImag = Math.sin(root);
  41.                 for (int i = 0; i < n; i += len * 2) {
  42.                     double wReal = 1;
  43.                     double wImag = 0;
  44.                     for (int j = 0; j < len; j++) {
  45.                         double real1 = real[i + j];
  46.                         double imag1 = imag[i + j];
  47.                         double real2 = real[i + j + len];
  48.                         double imag2 = imag[i + j + len];
  49.                         double real3 = real2 * wReal - imag2 * wImag;
  50.                         double imag3 = real2 * wImag + imag2 * wReal;
  51.                         real[i + j] = real1 + real3;
  52.                         imag[i + j] = imag1 + imag3;
  53.                         real[i + j + len] = real1 - real3;
  54.                         imag[i + j + len] = imag1 - imag3;
  55.                         double nextWReal = wReal * rootReal - wImag * rootImag;
  56.                         double nextWImag = wReal * rootImag + wImag * rootReal;
  57.                         wReal = nextWReal;
  58.                         wImag = nextWImag;
  59.                     }
  60.                 }
  61.             }
  62.  
  63.             if (inv) {
  64.                 for (int i = 0; i < n; i++) {
  65.                     real[i] /= n;
  66.                     imag[i] /= n;
  67.                 }
  68.             }
  69.         }
  70.  
  71.         void swap(double[] a, int i, int j) {
  72.             double t = a[i];
  73.             a[i] = a[j];
  74.             a[j] = t;
  75.         }
  76.  
  77.         double[] copy(long[] a, int n) {
  78.             double[] b = new double[n];
  79.             for (int i = 0; i < a.length; i++) b[i] = a[i];
  80.             return b;
  81.         }
  82.  
  83.         long[] multiply(long[] a, long[] b) {
  84.             int n = Math.max(1, Integer.highestOneBit(Math.max(a.length, b.length) - 1) << 2);
  85.             double[] aReal = copy(a, n);
  86.             double[] aImag = new double[n];
  87.             double[] bReal = copy(b, n);
  88.             double[] bImag = new double[n];
  89.             fft(aReal, aImag, false);
  90.             fft(bReal, bImag, false);
  91.             for (int i = 0; i < n; i++) {
  92.                 double real = aReal[i] * bReal[i] - aImag[i] * bImag[i];
  93.                 double imag = aReal[i] * bImag[i] + aImag[i] * bReal[i];
  94.                 aReal[i] = real;
  95.                 aImag[i] = imag;
  96.             }
  97.             fft(aReal, aImag, true);
  98.             long[] res = new long[n];
  99.             for (int i = 0; i < n; i++) res[i] = Math.round(aReal[i]);
  100.             return res;
  101.         }
  102.     }
Advertisement
Add Comment
Please, Sign In to add comment