Use montgomery in modint

Author: adamant-pwn

cp-algo/linalg/vector.hpp

@@ -127,9 +127,10 @@ namespace cp_algo::linalg {
         using Base::Base;
 
         void add_scaled(vec const& b, base scale, size_t i = 0) override {
+            uint64_t scaler = scale.getr();
             if(scale != base(0)) {
                 for(; i < size(*this); i++) {
-                    (*this)[i].add_unsafe(scale.getr() * b[i].getr());
+                    (*this)[i].add_unsafe(scaler * b[i].getr_direct());
                 }
                 if(++counter == 8) {
                     for(auto &it: *this) {

cp-algo/math/fft.hpp

@@ -1,7 +1,7 @@
 #ifndef CP_ALGO_MATH_FFT_HPP
 #define CP_ALGO_MATH_FFT_HPP
 #include "common.hpp"
-#include "cp-algo/number_theory/modint.hpp"
+#include "../number_theory/modint.hpp"
 #include <algorithm>
 #include <complex>
 #include <cassert>

cp-algo/math/poly.hpp

@@ -4,7 +4,7 @@
 #include "poly/impl/base.hpp"
 #include "poly/impl/div.hpp"
 #include "combinatorics.hpp"
-#include "cp-algo/number_theory/discrete_sqrt.hpp"
+#include "../number_theory/discrete_sqrt.hpp"
 #include "fft.hpp"
 #include <functional>
 #include <algorithm>

cp-algo/number_theory/discrete_log.hpp

@@ -26,9 +26,9 @@ namespace cp_algo::math {
             for(size_t k = 0; k < m; k += sqrtmod) {
                 auto it = small.find((base(c) * cur).getr());
                 if(it != end(small)) {
-                    auto cand = base::with_mod(period(base(b)), [&](){
-                        return base(it->second - k);
-                    }).getr();
+                    auto cand = base::with_mod(period(base(b)), [&]() {
+                        return base(it->second - k).getr();
+                    });
                     if(base(a) * bpow(base(b), cand) == base(c)) {
                         return cand;
                     } else {

cp-algo/number_theory/discrete_sqrt.hpp

@@ -1,8 +1,8 @@
 #ifndef CP_ALGO_NUMBER_THEORY_DISCRETE_SQRT_HPP
 #define CP_ALGO_NUMBER_THEORY_DISCRETE_SQRT_HPP
 #include "modint.hpp"
-#include "cp-algo/random/rng.hpp"
-#include "cp-algo/math/affine.hpp"
+#include "../random/rng.hpp"
+#include "../math/affine.hpp"
 namespace cp_algo::math {
     // https://en.wikipedia.org/wiki/Berlekamp-Rabin_algorithm
     template<modint_type base>

cp-algo/number_theory/factorize.hpp

@@ -1,7 +1,7 @@
 #ifndef CP_ALGO_MATH_FACTORIZE_HPP
 #define CP_ALGO_MATH_FACTORIZE_HPP
 #include "primality.hpp"
-#include "cp-algo/random/rng.hpp"
+#include "../random/rng.hpp"
 namespace cp_algo::math {
     // https://en.wikipedia.org/wiki/Pollard%27s_rho_algorithm
     void factorize(uint64_t m, std::vector<int64_t> &res) {

cp-algo/number_theory/modint.hpp

@@ -1,30 +1,63 @@
 #ifndef CP_ALGO_NUMBER_THEORY_MODINT_HPP
 #define CP_ALGO_NUMBER_THEORY_MODINT_HPP
-#include "cp-algo/math/common.hpp"
+#include "../math/common.hpp"
 #include <iostream>
+#include <cassert>
 namespace cp_algo::math {
+    inline constexpr uint64_t inv64(uint64_t x) {
+        assert(x % 2);
+        uint64_t y = 1;
+        while(y * x != 1) {
+            y *= 2 - x * y;
+        }
+        return y;
+    }
+
     template<typename modint>
     struct modint_base {
         static int64_t mod() {
             return modint::mod();
         }
+        static uint64_t imod() {
+            return modint::imod();
+        }
+        static __uint128_t pw128() {
+            return modint::pw128();
+        }
+        static uint64_t m_reduce(__uint128_t ab) {
+            if(mod() % 2 == 0) {
+                return ab % mod();
+            } else {
+                uint64_t m = ab * imod();
+                int64_t res = (ab + __uint128_t(m) * mod()) >> 64;
+                return res < mod() ? res : res - mod();
+            }
+        }
+        static uint64_t m_transform(uint64_t a) {
+            if(mod() % 2 == 0) {
+                return a;
+            } else {
+                return m_reduce(a * pw128());
+            }
+        }
         modint_base(): r(0) {}
         modint_base(int64_t rr): r(rr % mod()) {
             r = std::min(r, r + mod());
+            r = m_transform(r);
         }
         modint inv() const {
             return bpow(to_modint(), mod() - 2);
         }
-        modint operator - () const {return std::min(-r, mod() - r);}
+        modint operator - () const {
+            modint neg;
+            neg.r = std::min(-r, mod() - r);
+            return neg;
+        }
         modint& operator /= (const modint &t) {
             return to_modint() *= t.inv();
         }
         modint& operator *= (const modint &t) {
-            if(mod() <= uint32_t(-1)) {
-                r = r * t.r % mod();
-            } else {
-                r = __int128(r) * t.r % mod();
-            }
+            r = m_reduce(__uint128_t(r) * t.r);
             return to_modint();
         }
         modint& operator += (const modint &t) {
@@ -40,7 +73,10 @@ namespace cp_algo::math {
         modint operator * (const modint &t) const {return modint(to_modint()) *= t;}
         modint operator / (const modint &t) const {return modint(to_modint()) /= t;}
         auto operator <=> (const modint_base &t) const = default;
-        int64_t rem() const {return 2 * r > (uint64_t)mod() ? r - mod() : r;}
+        int64_t rem() const {
+            uint64_t R = getr();
+            return 2 * R > (uint64_t)mod() ? R - mod() : R;
+        }
 
         // Only use if you really know what you're doing!
         uint64_t modmod() const {return 8ULL * mod() * mod();};
@@ -52,16 +88,21 @@ namespace cp_algo::math {
             }
             return to_modint();
         }
-        uint64_t& setr() {return r;}
-        uint64_t getr() const {return r;}
+        void setr(uint64_t rr) {r = m_transform(rr);}
+        uint64_t getr() const {return m_reduce(r);}
+        void setr_direct(uint64_t rr) {r = rr;}
+        uint64_t getr_direct() const {return r;}
     private:
         uint64_t r;
         modint& to_modint() {return static_cast<modint&>(*this);}
         modint const& to_modint() const {return static_cast<modint const&>(*this);}
     };
     template<typename modint>
     std::istream& operator >> (std::istream &in, modint_base<modint> &x) {
-        return in >> x.setr();
+        uint64_t r;
+        auto &res = in >> r;
+        x.setr(r);
+        return res;
     }
     template<typename modint>
     std::ostream& operator << (std::ostream &out, modint_base<modint> const& x) {
@@ -73,14 +114,24 @@ namespace cp_algo::math {
 
     template<int64_t m>
     struct modint: modint_base<modint<m>> {
+        static constexpr uint64_t im = m % 2 ? inv64(-m) : 0;
+        static constexpr uint64_t r2 = __uint128_t(-1) % m + 1;
         static constexpr int64_t mod() {return m;}
+        static constexpr uint64_t imod() {return im;}
+        static constexpr __uint128_t pw128() {return r2;}
         using Base = modint_base<modint<m>>;
         using Base::Base;
     };
 
     struct dynamic_modint: modint_base<dynamic_modint> {
         static int64_t mod() {return m;}
-        static void switch_mod(int64_t nm) {m = nm;}
+        static uint64_t imod() {return im;}
+        static __uint128_t pw128() {return r2;}
+        static void switch_mod(int64_t nm) {
+            m = nm;
+            im = m % 2 ? inv64(-m) : 0;
+            r2 = __uint128_t(-1) % m + 1;
+        }
         using Base = modint_base<dynamic_modint>;
         using Base::Base;
 
@@ -95,7 +146,10 @@ namespace cp_algo::math {
         }
     private:
         static int64_t m;
+        static uint64_t im, r1, r2;
     };
-    int64_t dynamic_modint::m = 0;
+    int64_t dynamic_modint::m = 1;
+    uint64_t dynamic_modint::im = -1;
+    uint64_t dynamic_modint::r2 = 0;
 }
 #endif // CP_ALGO_MATH_MODINT_HPP

verify/linalg/inv.test.cpp

@@ -1,7 +1,6 @@
 // @brief Inverse Matrix
 #define PROBLEM "https://judge.yosupo.jp/problem/inverse_matrix"
 #pragma GCC optimize("Ofast,unroll-loops")
-#pragma GCC target("tune=native")
 #include "cp-algo/linalg/matrix.hpp"
 #include <bits/stdc++.h>
 
@@ -32,4 +31,4 @@ signed main() {
     while(t--) {
         solve();
     }
-}
+}

verify/linalg/prod.test.cpp

@@ -1,7 +1,6 @@
 // @brief Matrix Product
 #define PROBLEM "https://judge.yosupo.jp/problem/matrix_product"
 #pragma GCC optimize("Ofast,unroll-loops")
-#pragma GCC target("tune=native")
 #include "cp-algo/linalg/matrix.hpp"
 #include <bits/stdc++.h>
 
@@ -29,4 +28,4 @@ signed main() {
     while(t--) {
         solve();
     }
-}
+}

verify/number_theory/discrete_log.test.cpp

@@ -1,7 +1,6 @@
 // @brief Discrete Logarithm
 #define PROBLEM "https://judge.yosupo.jp/problem/discrete_logarithm_mod"
 #pragma GCC optimize("Ofast,unroll-loops")
-#pragma GCC target("tune=native")
 #include "cp-algo/number_theory/discrete_log.hpp"
 #include <bits/stdc++.h>
 
@@ -30,4 +29,4 @@ signed main() {
     while(t--) {
         solve();
     }
-}
+}