Introduce modint_type concept, fix fft

Author: adamant-pwn

cp-algo/algebra/fft.hpp

@@ -70,12 +70,12 @@ namespace cp_algo::algebra::fft {
         }
     }
 
-    template<int m>
+    template<modint_type base>
     struct dft {
         static constexpr int split = 1 << 15;
         std::vector<point> A;
 
-        dft(std::vector<modint<m>> const& a, size_t n): A(n) {
+        dft(std::vector<base> const& a, size_t n): A(n) {
             for(size_t i = 0; i < std::min(n, a.size()); i++) {
                 A[i] = point(
                     a[i].rem() % split,
@@ -91,7 +91,7 @@ namespace cp_algo::algebra::fft {
             assert(A.size() == B.A.size());
             size_t n = A.size();
             if(!n) {
-                return std::vector<modint<m>>();
+                return std::vector<base>();
             }
             std::vector<point> C(n), D(n);
             for(size_t i = 0; i < n; i++) {
@@ -103,11 +103,11 @@ namespace cp_algo::algebra::fft {
             reverse(begin(C) + 1, end(C));
             reverse(begin(D) + 1, end(D));
             int t = 2 * n;
-            std::vector<modint<m>> res(n);
+            std::vector<base> res(n);
             for(size_t i = 0; i < n; i++) {
-                modint<m> A0 = llround(C[i].real() / t);
-                modint<m> A1 = llround(C[i].imag() / t + D[i].imag() / t);
-                modint<m> A2 = llround(D[i].real() / t);
+                base A0 = llround(C[i].real() / t);
+                base A1 = llround(C[i].imag() / t + D[i].imag() / t);
+                base A2 = llround(D[i].real() / t);
                 res[i] = A0 + A1 * split - A2 * split * split;
             }
             return res;
@@ -128,18 +128,18 @@ namespace cp_algo::algebra::fft {
         return n;
     }
 
-    template<int m>
-    void mul(std::vector<modint<m>> &a, std::vector<modint<m>> b) {
+    template<modint_type base>
+    void mul(std::vector<base> &a, std::vector<base> const& b) {
         if(std::min(a.size(), b.size()) < magic) {
             mul_slow(a, b);
             return;
         }
         auto n = com_size(a.size(), b.size());
-        auto A = dft<m>(a, n);
+        auto A = dft<base>(a, n);
         if(a == b) {
             a = A * A;
         } else {
-            a = A * dft<m>(b, n);
+            a = A * dft<base>(b, n);
         }
     }
 }

cp-algo/algebra/modint.hpp

@@ -68,6 +68,9 @@ namespace cp_algo::algebra {
         return out << x.getr();
     }
 
+    template<typename modint>
+    concept modint_type = std::is_base_of_v<modint_base<modint>, modint>;
+
     template<int64_t m>
     struct modint: modint_base<modint<m>> {
         static constexpr int64_t mod() {return m;}

cp-algo/algebra/number_theory.hpp

@@ -6,8 +6,7 @@
 #include <optional>
 namespace cp_algo::algebra {
     // https://en.wikipedia.org/wiki/Berlekamp-Rabin_algorithm
-    template<typename base>
-    requires(std::is_base_of_v<modint_base<base>, base>)
+    template<modint_type base>
     std::optional<base> sqrt(base b) {
         if(b == base(0)) {
             return base(0);
@@ -28,8 +27,7 @@ namespace cp_algo::algebra {
         }
     }
 
-    template<typename base>
-    requires(std::is_base_of_v<modint_base<base>, base>)
+    template<modint_type base>
     bool is_prime_mod() {
         auto m = base::mod();
         if(m == 1 || m % 2 == 0) {