qwerty787788

FFT

Aug 7th, 2013
564
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Java 4.92 KB | None | 0 0
  1. import java.io.*;
  2. import java.util.*;
  3.  
  4. public class CF {
  5.  
  6.     void FFT(double[] re, double[] im, boolean invert) {
  7.         int n = re.length;
  8.         if (im.length != n)
  9.             throw new AssertionError("Sizes of arrays differ");
  10.         if (Integer.bitCount(n) != 1)
  11.             throw new AssertionError("N is not power of 2");
  12.         if (n != 1) {
  13.             int m = n / 2;
  14.             double[] re1 = new double[m];
  15.             double[] im1 = new double[m];
  16.             double[] re2 = new double[m];
  17.             double[] im2 = new double[m];
  18.             for (int i = 0; i < n; i += 2) {
  19.                 re1[i / 2] = re[i];
  20.                 im1[i / 2] = im[i];
  21.                 re2[i / 2] = re[i + 1];
  22.                 im2[i / 2] = im[i + 1];
  23.             }
  24.             FFT(re1, im1, invert);
  25.             FFT(re2, im2, invert);
  26.             double angle = (invert ? -1 : 1) * Math.PI * 2 / n;
  27.             double epsR = Math.cos(angle);
  28.             double epsI = Math.sin(angle);
  29.             double curR = 1;
  30.             double curI = 0;
  31.             for (int i = 0; i < m; i++) {
  32.                 double real = curR * re2[i] - curI * im2[i];
  33.                 double imag = curI * re2[i] + curR * im2[i];
  34.                 re[i] = re1[i] + real;
  35.                 im[i] = im1[i] + imag;
  36.                 real = -real;
  37.                 imag = -imag;
  38.                 re[i + m] = re1[i] + real;
  39.                 im[i + m] = im1[i] + imag;
  40.                 double nR = curR * epsR - curI * epsI;
  41.                 double nI = curR * epsI + curI * epsR;
  42.                 curR = nR;
  43.                 curI = nI;
  44.                 if (invert) {
  45.                     re[i] /= 2.;
  46.                     im[i] /= 2;
  47.                     re[i + m] /= 2;
  48.                     im[i + m] /= 2;
  49.                 }
  50.             }
  51.         }
  52.     }
  53.  
  54.     long[] mul(long[] a, long[] b) {
  55.         int len = Math.max(a.length, b.length) * 2;
  56.         int mLen = 1;
  57.         while (mLen < len)
  58.             mLen *= 2;
  59.         double[] r1 = new double[mLen];
  60.         double[] i1 = new double[mLen];
  61.         for (int i = 0; i < a.length; i++)
  62.             r1[i] = a[i];
  63.         double[] r2 = new double[mLen];
  64.         double[] i2 = new double[mLen];
  65.         for (int i = 0; i < b.length; i++)
  66.             r2[i] = b[i];
  67.         FFT(r1, i1, false);
  68.         FFT(r2, i2, false);
  69.         double[] rNew = new double[mLen];
  70.         double[] iNew = new double[mLen];
  71.         for (int i = 0; i < mLen; i++) {
  72.             rNew[i] = r1[i] * r2[i] - i1[i] * i2[i];
  73.             iNew[i] = r1[i] * i2[i] + r2[i] * i1[i];
  74.         }
  75.         FFT(rNew, iNew, true);
  76.         long[] res = new long[mLen];
  77.         for (int i = 0; i < mLen; i++)
  78.             res[i] = (long) Math.round(rNew[i]);
  79.         return res;
  80.     }
  81.  
  82.     long[] slowMul(long[] a, long[] b) {
  83.         long[] res = new long[a.length + b.length + 2];
  84.         for (int i = 0; i < a.length; i++)
  85.             for (int j = 0; j < b.length; j++)
  86.                 res[i + j] += a[i] * b[j];
  87.         return res;
  88.     }
  89.  
  90.     boolean same(long[] a, long[] b) {
  91.         for (int i = 0; i < Math.max(a.length, b.length); i++) {
  92.             if (i < a.length && i < b.length) {
  93.                 if (a[i] != b[i])
  94.                     return false;
  95.             } else {
  96.                 if (i < a.length) {
  97.                     if (a[i] != 0)
  98.                         return false;
  99.                 } else {
  100.                     if (b[i] != 0)
  101.                         return false;
  102.                 }
  103.             }
  104.         }
  105.         return true;
  106.     }
  107.  
  108.     void ACTest() {
  109.         for (int test = 0;; test++) {
  110.             System.err.println(test);
  111.             int n = 1;
  112.             long[] a = new long[n];
  113.             long[] b = new long[n];
  114.             for (int i = 0; i < n; i++) {
  115.                 a[i] = rnd.nextInt(100);
  116.                 b[i] = rnd.nextInt(100);
  117.             }
  118.             long[] r1 = mul(a, b);
  119.             long[] r2 = slowMul(a, b);
  120.             if (!same(r1, r2)) {
  121.                 throw new AssertionError();
  122.             }
  123.         }
  124.     }
  125.    
  126.     void TLTest() {
  127.         for (int test = 0;; test++) {
  128.             System.err.println(test);
  129.             long st = System.currentTimeMillis();
  130.             int n = 100000;
  131.             long[] a = new long[n];
  132.             long[] b = new long[n];
  133.             for (int i = 0; i < n; i++) {
  134.                 a[i] = rnd.nextInt(10000);
  135.                 b[i] = rnd.nextInt(10000);
  136.             }
  137.             long[] r1 = mul(a, b);
  138.             System.err.println(System.currentTimeMillis() - st + " ms");
  139.         }
  140.     }
  141.    
  142.     Random rnd = new Random(77);
  143.  
  144.     void realSolve() throws IOException {
  145.         TLTest();
  146.     }
  147.  
  148.     private class InputReader {
  149.         StringTokenizer st;
  150.         BufferedReader br;
  151.  
  152.         public InputReader(File f) {
  153.             try {
  154.                 br = new BufferedReader(new FileReader(f));
  155.             } catch (FileNotFoundException e) {
  156.                 e.printStackTrace();
  157.             }
  158.         }
  159.  
  160.         public InputReader(InputStream f) {
  161.             br = new BufferedReader(new InputStreamReader(f));
  162.         }
  163.  
  164.         String next() {
  165.             while (st == null || !st.hasMoreElements()) {
  166.                 String s;
  167.                 try {
  168.                     s = br.readLine();
  169.                 } catch (IOException e) {
  170.                     return null;
  171.                 }
  172.                 if (s == null)
  173.                     return null;
  174.                 st = new StringTokenizer(s);
  175.             }
  176.             return st.nextToken();
  177.         }
  178.  
  179.         int nextInt() {
  180.             return Integer.parseInt(next());
  181.         }
  182.  
  183.         double nextDouble() {
  184.             return Double.parseDouble(next());
  185.         }
  186.  
  187.         boolean hasMoreElements() {
  188.             while (st == null || !st.hasMoreElements()) {
  189.                 String s;
  190.                 try {
  191.                     s = br.readLine();
  192.                 } catch (IOException e) {
  193.                     return false;
  194.                 }
  195.                 st = new StringTokenizer(s);
  196.             }
  197.             return st.hasMoreElements();
  198.         }
  199.  
  200.         long nextLong() {
  201.             return Long.parseLong(next());
  202.         }
  203.     }
  204.  
  205.     InputReader in;
  206.     PrintWriter out;
  207.  
  208.     void solveIO() throws IOException {
  209.         in = new InputReader(System.in);
  210.         out = new PrintWriter(System.out);
  211.  
  212.         realSolve();
  213.  
  214.         out.close();
  215.  
  216.     }
  217.  
  218.     public static void main(String[] args) throws IOException {
  219.         new CF().solveIO();
  220.     }
  221. }
Advertisement
Add Comment
Please, Sign In to add comment