fix a bug in fft::trans

zimpha · zimpha · commit 202860b0b2b9 · 2017-08-24T00:14:36.000+08:00
diff --git a/mathematics/fast-fourier-transform.cc b/mathematics/fast-fourier-transform.cc
@@ -15,6 +15,12 @@ namespace fft {
     Complex operator - (const Complex& r) const {
       return Complex(x - r.x , y - r.y);
     }
+    Complex operator * (const double k) const {
+      return Complex(x * k, y * k);
+    }
+    Complex operator / (const double k) const {
+      return Complex(x / k, y / k);
+    }
     Complex operator * (const Complex& r) const {
       return Complex(x * r.x - y * r.y , x * r.y + y * r.x);
     }
@@ -54,9 +60,14 @@ namespace fft {
         }
       }
     }
+    if (oper == -1) {
+      for (int i = 0; i < n; ++i) {
+        P[i] = P[i] / n;
+      }
+    }
   }
   Complex A[N] , B[N] , C1[N] , C2[N];
-  std::vector<int> conv(const std::vector<int> &a, const std::vector<int> &b) {
+  std::vector<int64> conv(const std::vector<int> &a, const std::vector<int> &b) {
     int n = a.size(), m = b.size(), s = 1;
     while (s <= n + m - 2) s <<= 1;
     init(s);
@@ -69,8 +80,11 @@ namespace fft {
     for (int i = 0; i < s; ++i) {
       A[i] = A[i] * B[i];
     }
+    for (int i = 0; i < s; ++i) {
+      w[i] = w[i].conj();
+    }
     trans(A, s, -1);
-    std::vector<int> res(n + m - 1);
+    std::vector<int64> res(n + m - 1);
     for (int i = 0; i < s; ++i) {
       res[i] = (int64)(A[i].x + 0.5);
     }
@@ -101,10 +115,10 @@ namespace fft {
     trans(C1, s, -1);
     trans(C2, s, -1);
     for (int i = 0 ; i < n + m - 1; ++i) {
-      int x = (int64)(C1[i].x / s + 0.5) % mod;
-      int y1 = (int64)(C1[i].y / s + 0.5) % mod;
-      int y2 = (int64)(C2[i].x / s + 0.5) % mod;
-      int z = (int64)(C2[i].y / s + 0.5) % mod;
+      int x = int64(C1[i].x + 0.5) % mod;
+      int y1 = int64(C1[i].y + 0.5) % mod;
+      int y2 = int64(C2[i].x + 0.5) % mod;
+      int z = int64(C2[i].y + 0.5) % mod;
       res[i] = ((int64)x * M * M + (int64)(y1 + y2) * M + z) % mod;
     }
   }

Original file line number	Diff line number	Diff line change
`@@ -15,6 +15,12 @@ namespace fft {`
`15`	`15`	`Complex operator - (const Complex& r) const {`
`16`	`16`	`return Complex(x - r.x , y - r.y);`
`17`	`17`	`}`
	`18`	`+ Complex operator * (const double k) const {`
	`19`	`+ return Complex(x * k, y * k);`
	`20`	`+ }`
	`21`	`+ Complex operator / (const double k) const {`
	`22`	`+ return Complex(x / k, y / k);`
	`23`	`+ }`
`18`	`24`	`Complex operator * (const Complex& r) const {`
`19`	`25`	`return Complex(x * r.x - y * r.y , x * r.y + y * r.x);`
`20`	`26`	`}`
`@@ -54,9 +60,14 @@ namespace fft {`
`54`	`60`	`}`
`55`	`61`	`}`
`56`	`62`	`}`
	`63`	`+ if (oper == -1) {`
	`64`	`+ for (int i = 0; i < n; ++i) {`
	`65`	`+ P[i] = P[i] / n;`
	`66`	`+ }`
	`67`	`+ }`
`57`	`68`	`}`
`58`	`69`	`Complex A[N] , B[N] , C1[N] , C2[N];`
`59`		`- std::vector<int> conv(const std::vector<int> &a, const std::vector<int> &b) {`
	`70`	`+ std::vector<int64> conv(const std::vector<int> &a, const std::vector<int> &b) {`
`60`	`71`	`int n = a.size(), m = b.size(), s = 1;`
`61`	`72`	`while (s <= n + m - 2) s <<= 1;`
`62`	`73`	`init(s);`
`@@ -69,8 +80,11 @@ namespace fft {`
`69`	`80`	`for (int i = 0; i < s; ++i) {`
`70`	`81`	`A[i] = A[i] * B[i];`
`71`	`82`	`}`
	`83`	`+ for (int i = 0; i < s; ++i) {`
	`84`	`+ w[i] = w[i].conj();`
	`85`	`+ }`
`72`	`86`	`trans(A, s, -1);`
`73`		`- std::vector<int> res(n + m - 1);`
	`87`	`+ std::vector<int64> res(n + m - 1);`
`74`	`88`	`for (int i = 0; i < s; ++i) {`
`75`	`89`	`res[i] = (int64)(A[i].x + 0.5);`
`76`	`90`	`}`
`@@ -101,10 +115,10 @@ namespace fft {`
`101`	`115`	`trans(C1, s, -1);`
`102`	`116`	`trans(C2, s, -1);`
`103`	`117`	`for (int i = 0 ; i < n + m - 1; ++i) {`
`104`		`- int x = (int64)(C1[i].x / s + 0.5) % mod;`
`105`		`- int y1 = (int64)(C1[i].y / s + 0.5) % mod;`
`106`		`- int y2 = (int64)(C2[i].x / s + 0.5) % mod;`
`107`		`- int z = (int64)(C2[i].y / s + 0.5) % mod;`
	`118`	`+ int x = int64(C1[i].x + 0.5) % mod;`
	`119`	`+ int y1 = int64(C1[i].y + 0.5) % mod;`
	`120`	`+ int y2 = int64(C2[i].x + 0.5) % mod;`
	`121`	`+ int z = int64(C2[i].y + 0.5) % mod;`
`108`	`122`	`res[i] = ((int64)x * M * M + (int64)(y1 + y2) * M + z) % mod;`
`109`	`123`	`}`
`110`	`124`	`}`