add fast miller-robin and prime factorization

zimpha · zimpha · commit c41078495881 · 2018-04-30T16:22:57.000+08:00
diff --git a/mathematics/Fibonacci.cc b/mathematics/Fibonacci.cc
@@ -1,16 +1,15 @@
 #include "basic.hpp"
+#include <functional>
 
-void fib(int64 n, int64 &x, int64 &y) {// store in x, n-th
-  if (n == 1) {
-    x = y = 1;
-    return;
-  } else if (n & 1) {
-    fib(n - 1, y, x);
-    y += x;
+// f(0) = 0, f(1) = 1, f(n) = f(n - 1) + f(n - 2)
+std::pair<int64, int64> fib(int64 n, int64 mod) {
+  if (n == 0) return {0, 1};
+  int64 x, y;
+  if (n & 1) {
+    std::tie(y, x) = fib(n - 1, mod);
+    return {x, (y + x) % mod};
   } else {
-    int64 a, b;
-    fib(n >> 1, a, b);
-    y = a * a + b * b;
-    x = a * b + a * (b - a);
+    std::tie(x, y) = fib(n >> 1, mod);
+    return {(x * y + x * (y - x + mod)) % mod, (x * x + y * y) % mod};
   }
 }
diff --git a/mathematics/ModInt.hpp b/mathematics/ModInt.hpp
@@ -12,6 +12,12 @@ struct UnsafeMod {
   UnsafeMod(): x(0) {}
   UnsafeMod(word _x): x(init(_x)) {}
 
+  bool operator == (const UnsafeMod& rhs) const {
+    return x == rhs.x;
+  }
+  bool operator != (const UnsafeMod& rhs) const {
+    return x != rhs.x;
+  }
   UnsafeMod& operator += (const UnsafeMod& rhs) {
     if ((x += rhs.x) >= mod) x -= mod;
     return *this;
diff --git a/mathematics/prime.cc b/mathematics/prime.cc
@@ -1,56 +1,220 @@
-#include "basic.hpp"
+#include <cmath>
+#include <numeric>
+#include <cassert>
+#include <cstdio>
 #include <vector>
 #include <cstdlib>
 #include <algorithm>
 
-struct Primality {
-public:
-  // 用miller rabin素数测试判断n是否为质数
-  bool is_prime(int64 n) {
-    if (n <= 1) return false;
-    if (n <= 3) return true;
-    if (~n & 1) return false;
-    const int u[] = {2,3,5,7,325,9375,28178,450775,9780504,1795265022,0};
-    int64 e = n - 1, a, c = 0; // 原理：http://miller-rabin.appspot.com/
-    while (~e & 1) e >>= 1, ++c;
-    for (int i = 0; u[i]; ++i) {
-      if (n <= u[i]) return true;
-      a = pow_mod(u[i], e, n);
-      if (a == 1) continue;
-      for (int j = 1; a != n - 1; ++j) {
-        if (j == c) return false;
-        a = mul_mod(a, a, n);
+namespace prime {
+
+using uint128 = __uint128_t;
+using uint64 = unsigned long long;
+using int64 = long long;
+using uint32 = unsigned int;
+using pii = std::pair<uint64, uint32>;
+
+inline uint64 sqr(uint64 x) { return x * x; }
+inline uint32 isqrt(uint64 x) { return sqrtl(x); }
+inline uint32 ctz(uint64 x) { return __builtin_ctzll(x); }
+
+template <typename word>
+word gcd(word a, word b) {
+  while (b) { word t = a % b; a = b; b = t; }
+  return a;
+}
+
+template <typename word, typename dword, typename sword>
+struct Mod {
+  Mod(): x(0) {}
+  Mod(word _x): x(init(_x)) {}
+  bool operator == (const Mod& rhs) const { return x == rhs.x; }
+  bool operator != (const Mod& rhs) const { return x != rhs.x; }
+  Mod& operator += (const Mod& rhs) { if ((x += rhs.x) >= mod) x -= mod; return *this; }
+  Mod& operator -= (const Mod& rhs) { if (sword(x -= rhs.x) < 0) x += mod; return *this; }
+  Mod& operator *= (const Mod& rhs) { x = reduce(dword(x) * rhs.x); return *this; }
+  Mod operator + (const Mod &rhs) const { return Mod(*this) += rhs; }
+  Mod operator - (const Mod &rhs) const { return Mod(*this) -= rhs; }
+  Mod operator * (const Mod &rhs) const { return Mod(*this) *= rhs; }
+  Mod operator - () const { return Mod() - *this; }
+  Mod pow(uint64 e) const {
+    Mod ret(1);
+    for (Mod base = *this; e; e >>= 1, base *= base) {
+      if (e & 1) ret *= base;
+    }
+    return ret;
+  }
+  word get() const { return reduce(x); }
+  static constexpr int word_bits = sizeof(word) * 8;
+  static word modulus() { return mod; }
+  static word init(word w) { return reduce(dword(w) * r2); }
+  static void set_mod(word m) { mod = m, inv = mul_inv(mod), r2 = -dword(mod) % mod; }
+  static word reduce(dword x) {
+    word y = word(x >> word_bits) - word((dword(word(x) * inv) * mod) >> word_bits);
+    return sword(y) < 0 ? y + mod : y;
+  }
+  static word mul_inv(word n, int e = 6, word x = 1) {
+    return !e ? x : mul_inv(n, e - 1, x * (2 - x * n));
+  }
+  static word mod, inv, r2;
+
+  word x;
+};
+
+using Mod64 = Mod<uint64, uint128, int64>;
+using Mod32 = Mod<uint32, uint64, int>;
+template <> uint64 Mod64::mod = 0;
+template <> uint64 Mod64::inv = 0;
+template <> uint64 Mod64::r2 = 0;
+template <> uint32 Mod32::mod = 0;
+template <> uint32 Mod32::inv = 0;
+template <> uint32 Mod32::r2 = 0;
+
+template <class word, class mod>
+bool composite(word n, const uint32* bases, int m) {
+  mod::set_mod(n);
+  int s = __builtin_ctzll(n - 1);
+  word d = (n - 1) >> s;
+  mod one{1}, minus_one{n - 1};
+  for (int i = 0, j; i < m; ++i) {
+    mod a = mod(bases[i]).pow(d);
+    if (a == one || a == minus_one) continue;
+    for (j = s - 1; j > 0; --j) {
+      if ((a *= a) == minus_one) break;
+    }
+    if (j == 0) return true;
+  }
+  return false;
+}
+
+bool is_prime(uint64 n) { // reference: http://miller-rabin.appspot.com
+  assert(n < (uint64(1) << 63));
+  static const uint32 bases[][7] = {
+    {2, 3},
+    {2, 299417},
+    {2, 7, 61},
+    {15, 176006322, uint32(4221622697)},
+    {2, 2570940, 211991001, uint32(3749873356)},
+    {2, 2570940, 880937, 610386380, uint32(4130785767)},
+    {2, 325, 9375, 28178, 450775, 9780504, 1795265022}
+  };
+  if (n <= 1) return false;
+  if (!(n & 1)) return n == 2;
+  if (n <= 8) return true;
+  int x = 6, y = 7;
+  if (n < 1373653) x = 0, y = 2;
+  else if (n < 19471033) x = 1, y = 2;
+  else if (n < 4759123141) x = 2, y = 3;
+  else if (n < 154639673381) x = y = 3;
+  else if (n < 47636622961201) x = y = 4;
+  else if (n < 3770579582154547) x = y = 5;
+  if (n < (uint32(1) << 31)) {
+    return !composite<uint32, Mod32>(n, bases[x], y);
+  } else if (n < (uint64(1) << 63)) {
+    return !composite<uint64, Mod64>(n, bases[x], y);
+  }
+  return true;
+}
+
+struct ExactDiv {
+  ExactDiv() {}
+  ExactDiv(uint64 n) : n(n), i(Mod64::mul_inv(n)), t(uint64(-1) / n) {}
+  friend uint64 operator / (uint64 n, ExactDiv d) { return n * d.i; };
+  bool divide(uint64 n) { return n / *this <= t; }
+  uint64 n, i, t;
+};
+
+std::vector<ExactDiv> primes;
+
+void init(uint32 n) {
+  uint32 sqrt_n = sqrt(n);
+  std::vector<bool> is_prime(n + 1, 1);
+  primes.clear();
+  for (uint32 i = 2; i <= sqrt_n; ++i) if (is_prime[i]) {
+    if (i != 2) primes.push_back(ExactDiv(i));
+    for (uint32 j = i * i; j <= n; j += i) is_prime[j] =  0;
+  }
+}
+
+template <typename word, typename mod>
+word brent(word n, word c) { // n must be composite and odd.
+  const uint64 s = 256;
+  mod::set_mod(n);
+  const mod one = mod(1), mc = mod(c);
+  mod y = one;
+  for (uint64 l = 1; ; l <<= 1) {
+    auto x = y;
+    for (int i = 0; i < (int)l; ++i) y = y * y + mc;
+    mod p = one;
+    for (int k = 0; k < (int)l; k += s) {
+      auto sy = y;
+      for (int i = 0; i < (int)std::min(s, l - k); ++i) {
+        y = y * y + mc;
+        p *= y - x;
+      }
+      word g = gcd(n, p.x);
+      if (g == 1) continue;
+      if (g == n) for (g = 1, y = sy; g == 1; ) {
+        y = y * y + mc, g = gcd(n, (y - x).x);
       }
+      return g;
     }
-    return true;
-  }
-  // 求一个小于n的因数，期望复杂度为O(n^0.25)，当n为非合数时返回n本身
-  int64 pollard_rho(int64 n){
-    if (n <= 3 || is_prime(n)) return n; // 保证n为合数时可去掉这行
-    while (1) {
-      int i = 1, cnt = 2;
-      int64 x = rand() % n, y = x, c = rand() % n;
-      if (!c || c == n - 2) ++c;
-      do {
-        int64 u = gcd(n - x + y, n);
-        if (u > 1 && u < n) return u;
-        if (++i == cnt) y = x, cnt <<= 1;
-        x = (c + mul_mod(x, x, n)) % n;
-      } while (x != y);
+  }
+}
+
+uint64 brent(uint64 n, uint64 c) {
+  if (n < (uint32(1) << 31)) {
+    return brent<uint32, Mod32>(n, c);
+  } else if (n < (uint64(1) << 63)) {
+    return brent<uint64, Mod64>(n, c);
+  }
+  return 0;
+}
+
+std::vector<pii> factors(uint64 n) {
+  assert(n < (uint64(1) << 63));
+  if (n <= 1) return {};
+  std::vector<pii> ret;
+  uint32 v = sqrtl(n);
+  if (uint64(v) * v == n) {
+    ret = factors(v);
+    for (auto &&e: ret) e.second *= 2;
+    return ret;
+  }
+  v = cbrtl(n);
+  if (uint64(v) * v * v == n) {
+    ret = factors(v);
+    for (auto &&e: ret) e.second *= 3;
+    return ret;
+  }
+  if (!(n & 1)) {
+    uint32 e = __builtin_ctzll(n);
+    ret.emplace_back(2, e);
+    n >>= e;
+  }
+  uint64 lim = sqr(primes.back().n);
+  for (auto &&p: primes) {
+    if (sqr(p.n) > n) break;
+    if (p.divide(n)) {
+      uint32 e = 1; n = n / p;
+      while (p.divide(n)) n = n / p, e++;
+      ret.emplace_back(p.n, e);
     }
-    return n;
-  }
-  // 使用rho方法对n做质因数分解，建议先筛去小质因数后再用此函数
-  std::vector<int64> factorize(int64 n){
-    std::vector<int64> u;
-    if (n > 1) u.push_back(n);
-    for (size_t i = 0; i < u.size(); ++i){
-      int64 x = pollard_rho(u[i]);
-      if(x == u[i]) continue;
-      u[i--] /= x;
-      u.push_back(x);
+  }
+
+  uint32 s = ret.size();
+  while (n > lim && !is_prime(n)) {
+    for (uint64 c = 1; ; ++c) {
+      uint64 p = brent(n, c);
+      if (!is_prime(p)) continue;
+      uint32 e = 1; n /= p;
+      while (n % p == 0) n /= p, e += 1;
+      ret.emplace_back(p, e);
+      break;
     }
-    std::sort(u.begin(), u.end());
-    return u;
   }
-};
+  if (n > 1) ret.emplace_back(n, 1);
+  if (ret.size() - s >= 2) sort(ret.begin() + s, ret.end());
+  return ret;
+}
+}
