Alex's Anthology of Algorithms Common Code for Contests in Concise C++
Mathematics / Number Theory

6.3.9 Prime Tests and Factorization

6-Mathematics/6.3.9_Prime_Tests_and_Factorization.cpp

Provides primality tests, prime factorization, and divisor generation. Alongside generic trial division, it includes deterministic Miller-Rabin and Pollard's rho for signed 64-bit integers. Prime factorizations are represented as sorted vectors of (prime, exponent) pairs. For $0$ and $1$, the prime factorization is empty.

Trial division tests primality and computes prime factorizations by trying possible divisors through the square root.

  • is_prime_slow(n) returns whether the integer n is prime using trial division.
  • factorize_slow(n) returns the prime factorization of n using trial division.
Time Complexity
$O(\sqrt{n})$ per call to is_prime_slow() and factorize_slow().
Space Complexity
$O(1)$ auxiliary for is_prime_slow() and $O(f)$ for the returned factorization, where $f$ is the number of compressed prime factors.

Implementation

#include <algorithm>
#include <cassert>
#include <chrono>
#include <cstdint>
#include <numeric>
#include <random>
#include <utility>
#include <vector>

template<typename Int>
bool is_prime_slow(Int n) {
  if (n == 2 || n == 3) {
    return true;
  }
  if (n < 2 || n % 2 == 0 || n % 3 == 0) {
    return false;
  }
  for (Int i = 5, w = 4; i <= n / i; i += w) {
    if (n % i == 0) {
      return false;
    }
    w = 6 - w;
  }
  return true;
}

template<typename Int>
std::vector<std::pair<Int, int>> factorize_slow(Int n) {
  if (n <= 1) {
    return {};
  }
  std::vector<std::pair<Int, int>> res;
  for (Int i = 2; i <= n / i; i++) {
    if (n % i == 0) {
      int exponent = 0;
      do {
        n /= i;
        exponent++;
      } while (n % i == 0);
      res.emplace_back(i, exponent);
    }
  }
  if (n > 1) {
    res.emplace_back(n, 1);
  }
  return res;
}

The Miller-Rabin test writes $n - 1 = d \cdot 2^r$ and repeatedly squares $a^d$; a base that breaks the pattern required of a prime proves that $n$ is composite.

  • is_probable_prime(n, k = 20) returns whether n is probably prime using k random Miller-Rabin bases. The result is guaranteed correct when n is prime. For composite n, a true result has one-sided error probability at most $(1 / 4)^k$.
  • is_prime(n) returns whether the signed 64-bit integer n is prime using deterministic Miller-Rabin bases sufficient up to and including $2^{63} - 1$.

Modular multiplication uses __uint128_t when available, with a slower portable double-and-add fallback that avoids overflow on compilers without 128-bit integers.

Time Complexity
$O(k \cdot \log n)$ modular multiplications per call to is_probable_prime() and $O(\log n)$ per call to is_prime(). The portable multiplication fallback adds another $O(\log n)$ factor.
Space Complexity
$O(1)$ auxiliary.

Implementation

uint64_t mulmod(uint64_t a, uint64_t b, uint64_t m) {
  assert(m > 0);
#if defined(__SIZEOF_INT128__)
  return static_cast<uint64_t>(static_cast<__uint128_t>(a) * b % m);
#else
  uint64_t res = 0, cur = a % m;
  for (; b > 0; b >>= 1) {
    if (b & 1) {
      res = res >= m - cur ? res - (m - cur) : res + cur;
    }
    cur = cur >= m - cur ? cur - (m - cur) : cur + cur;
  }
  return res;
#endif
}

uint64_t powmod(uint64_t x, uint64_t n, uint64_t m) {
  assert(m > 0);
  uint64_t res = 1 % m;
  for (; n > 0; n >>= 1) {
    if (n & 1) {
      res = mulmod(res, x, m);
    }
    x = mulmod(x, x, m);
  }
  return res;
}

uint64_t rand64u() {
  static std::mt19937_64 rng(std::chrono::steady_clock::now().time_since_epoch().count());
  return rng();
}

bool is_probable_prime(int64_t n, int k = 20) {
  if (n == 2 || n == 3) {
    return true;
  }
  if (n < 2 || n % 2 == 0 || n % 3 == 0) {
    return false;
  }
  uint64_t s = n - 1, p = n - 1;
  while (!(s & 1)) {
    s >>= 1;
  }
  for (int i = 0; i < k; i++) {
    uint64_t x, r = powmod(rand64u() % (n - 3) + 2, s, n);  // random witness base in [2, n - 2]
    for (x = s; x != p && r != 1 && r != p; x <<= 1) {
      r = mulmod(r, r, n);
    }
    if (r != p && !(x & 1)) {
      return false;
    }
  }
  return true;
}

bool is_prime(int64_t n) {
  if (n < 2) {
    return false;
  }
  static const int small_primes[] = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29};
  for (int p : small_primes) {
    if (n % p == 0) {
      return n == p;
    }
  }
  if (n < 31 * 31) {
    return true;
  }
  uint64_t t;
  int s = 0;
  for (t = n - 1; !(t & 1); t >>= 1) {
    s++;
  }
  static const uint64_t bases[] = {2, 325, 9375, 28178, 450775, 9780504, 1795265022};
  for (uint64_t a : bases) {
    if (a % n == 0) {
      continue;
    }
    uint64_t r = powmod(a, t, n);
    if (r == 1) {
      continue;
    }
    bool ok = false;
    for (int j = 0; j < s && !ok; j++) {
      ok |= (r == static_cast<uint64_t>(n) - 1);
      r = mulmod(r, r, n);
    }
    if (!ok) {
      return false;
    }
  }
  return true;
}

Pollard's rho iterates a pseudorandom map modulo $n$; two iterates that collide modulo a hidden prime factor reveal that factor through a GCD with their difference. The complete factorization routine first removes small prime factors with a cached linear sieve, then uses Miller-Rabin and Pollard's rho on the remaining cofactor. Factorizations remain sorted and compressed as (prime, exponent) pairs.

  • rho_factor(n) returns a factor of n that is not necessarily prime using Pollard's rho with Brent's optimization. If n is prime, then n itself is returned. The algorithm may need to be retried to find a nontrivial factor, as done in factorize_rho().
  • factorize_rho(n) returns the prime factorization of a 64-bit integer using Miller-Rabin and Pollard's rho without initial trial division.
  • factorize(n, small_prime_limit = 1000000) returns the prime factorization of a 64-bit integer n, first testing primes through small_prime_limit and then falling back to Pollard's rho. It supports integers up to and including $2^{63} - 1$.
  • divisors_from_factors(factors) returns all divisors from a compressed factorization.
  • divisors(n) returns all divisors of a 64-bit integer using factorize() followed by divisors_from_factors().
Time Complexity
  • $O(\sqrt{p})$ expected modular multiplications per call to rho_factor(n), where $p$ is the smallest prime factor of n. For composite n, this is at most $O(n^{1/4})$.
  • $O(n^{1/4})$ expected per call to factorize_rho(). factorize() additionally performs $O(L/\log L)$ trial divisions, where $L$ is small_prime_limit, and takes $O(L)$ time if the cache is rebuilt.
  • $O(d \log d)$ per call to divisors_from_factors(), where $d$ is the number of divisors generated.
  • $O(L/\log L + n^{1/4} + d \log d)$ expected per call to divisors(), plus $O(L)$ if the cache is rebuilt.
Space Complexity
  • $O(c)$ cached space after sieving through $c$.
  • $O(f + \log n)$ auxiliary for recursive factorization, where $f$ is the number of compressed prime factors returned.
  • $O(d)$ output space and $O(\log d)$ auxiliary stack space for divisors_from_factors() and divisors().

Implementation

using Factors = std::vector<std::pair<int64_t, int>>;

int64_t rho_factor(int64_t n) {
  if (n % 2 == 0) {
    return 2;
  }
  uint64_t y = rand64u() % (n - 1) + 1;
  uint64_t c = rand64u() % (n - 1) + 1;
  uint64_t m = rand64u() % (n - 1) + 1;
  uint64_t g = 1, r = 1, q = 1, ys = 0, x = 0;
  for (r = 1; g == 1; r <<= 1) {
    x = y;
    for (uint64_t i = 0; i < r; i++) {
      y = (mulmod(y, y, n) + c) % n;
    }
    for (uint64_t k = 0; k < r && g == 1; k += m) {
      ys = y;
      int64_t lim = std::min(m, r - k);
      for (int j = 0; j < lim; j++) {
        y = (mulmod(y, y, n) + c) % n;
        q = mulmod(q, (x > y) ? (x - y) : (y - x), n);
      }
      g = std::gcd(q, n);
    }
  }
  if (g == static_cast<uint64_t>(n)) {
    do {
      ys = (mulmod(ys, ys, n) + c) % n;
      g = std::gcd((x > ys) ? (x - ys) : (ys - x), n);
    } while (g <= 1);
  }
  return g;
}

Factors factorize_rho(int64_t n) {
  std::vector<int64_t> factors;
  auto collect = [&](auto &&collect, int64_t value) {
    if (value <= 1) {
      return;
    }
    if (is_prime(value)) {
      factors.push_back(value);
      return;
    }
    int64_t p;
    do {
      p = rho_factor(value);
    } while (p == value);
    collect(collect, p);
    collect(collect, value / p);
  };
  collect(collect, n);
  std::sort(factors.begin(), factors.end());
  Factors res;
  for (int64_t p : factors) {
    if (res.empty() || res.back().first != p) {
      res.emplace_back(p, 1);
    } else {
      res.back().second++;
    }
  }
  return res;
}

static const auto &cached_sieve(int n) {
  struct Cache {
    int limit = 1;
    std::vector<int> least{0, 1}, primes;
  };
  static Cache cache;
  if (n <= cache.limit) {
    return cache;
  }
  cache.limit = std::max(1, n);
  cache.least.assign(cache.limit + 1, 0);
  cache.primes.clear();
  for (int i = 2; i <= cache.limit; i++) {
    if (cache.least[i] == 0) {
      cache.least[i] = i;
      cache.primes.push_back(i);
    }
    for (int p : cache.primes) {
      if (p > cache.least[i] || i > cache.limit / p) {
        break;
      }
      cache.least[i * p] = p;
    }
  }
  return cache;
}

Factors factorize(int64_t n, int small_prime_limit = 1000000) {
  assert(small_prime_limit >= 1);
  if (n <= 1) {
    return {};
  }
  const auto &sieve = cached_sieve(small_prime_limit);
  if (n <= small_prime_limit) {
    Factors res;
    while (n > 1) {
      int p = sieve.least[n], cnt = 0;
      do {
        n /= p;
        cnt++;
      } while (n > 1 && sieve.least[n] == p);
      res.emplace_back(p, cnt);
    }
    return res;
  }
  Factors res;
  const std::vector<int> &primes = sieve.primes;
  for (int p : primes) {
    if (p > small_prime_limit || static_cast<int64_t>(p) > n / p) {
      break;
    }
    if (n % p == 0) {
      int cnt = 0;
      do {
        n /= p;
        cnt++;
      } while (n % p == 0);
      res.emplace_back(p, cnt);
    }
  }
  Factors tail = factorize_rho(n);
  res.insert(res.end(), tail.begin(), tail.end());
  return res;
}

std::vector<int64_t> divisors_from_factors(const Factors &factors) {
  std::vector<int64_t> res{1};
  for (const auto &factor : factors) {
    int old_size = static_cast<int>(res.size());
    int64_t power = 1;
    for (int e = 1; e <= factor.second; e++) {
      power *= factor.first;
      for (int i = 0; i < old_size; i++) {
        res.push_back(res[i] * power);
      }
    }
  }
  std::sort(res.begin(), res.end());
  return res;
}

std::vector<int64_t> divisors(int64_t n) {
  if (n <= 1) {
    return (n < 1) ? std::vector<int64_t>() : std::vector<int64_t>(1, 1);
  }
  return divisors_from_factors(factorize(n));
}

Example Usage

#include <ctime>
#include <iomanip>
#include <iostream>
#include <set>
using namespace std;

void validate(int64_t n, const Factors &factors) {
  if (n <= 1) {
    assert(factors.empty());
    return;
  }
  if (is_prime(n)) {
    assert((factors == Factors{{n, 1}}));
    return;
  }
  int64_t prod = 1;
  for (const auto &factor : factors) {
    assert(is_prime(factor.first));
    for (int i = 0; i < factor.second; i++) {
      prod *= factor.first;
    }
  }
  assert(prod == n);
}

int main() {
  {  // Small primality tests.
    vector<pair<int, bool>> cases{
        {-10, false}, {-1, false}, {0, false}, {1, false}, // non-primes
        {2, true}, {3, true}, {4, false}, {5, true},
        {9, false},    // divisible by 3
        {25, false},   // square of small prime
        {29, true},    // last small-prime table entry
        {31, true},    // just past small-prime table
        {49, false},   // square, not divisible by 2 or 3
        {121, false},  // 11^2
        {169, false},  // 13^2
        {961, false},  // 31^2, boundary for n < 31*31 shortcut
        {997, true},   // prime near 1000
        {1001, false}, // 7*11*13
        {2047, false}, // 23*89, pseudoprime to base 2
        {341, false},  // 11*31, Fermat pseudoprime to base 2
        {561, false},  // Carmichael number
    };
    for (const auto &[n, expected] : cases) {
      assert(is_prime_slow(n) == expected);
      assert(is_probable_prime(n) == expected);
      assert(is_prime(n) == expected);
    }
  }
  {  // Primality test benchmark.
    vector<int64_t> nums{
        772023803LL, 2147483647LL, 5705234089LL, 6339503641LL, 999966000289LL,
    };
    cout << "Primality test:" << endl;
    cout << setw(15) << left << "num:";
    cout << setw(18) << left << "is_prime_slow():";
    cout << setw(20) << left << "is_probable_prime():";
    cout << setw(13) << right << "is_prime():" << endl;
    cout << fixed << setprecision(3);
    for (int64_t n : nums) {
      clock_t start = clock();
      bool p1 = is_prime_slow(n);
      double t1 = static_cast<double>(clock() - start) / CLOCKS_PER_SEC;
      start = clock();
      bool p2 = is_probable_prime(n);
      double t2 = static_cast<double>(clock() - start) / CLOCKS_PER_SEC;
      start = clock();
      bool p3 = is_prime(n);
      double t3 = static_cast<double>(clock() - start) / CLOCKS_PER_SEC;
      assert(p1 == p2 && p1 == p3);
      cout << setw(12) << left << n << right;
      cout << setw(14) << 1000 * t1 << "ms";
      cout << setw(16) << 1000 * t2 << "ms";
      cout << setw(16) << 1000 * t3 << "ms" << endl;
    }
    cout << endl;
  }
  {  // Small factorization tests.
    for (int64_t i = 1; i <= 10000; i++) {
      auto v1 = factorize_slow(i);
      auto v2 = factorize(i);
      validate(i, v1);
      assert(v1 == v2);
      auto d = divisors(i);
      set<int> s(d.begin(), d.end());
      assert(d.size() == s.size());
      for (int j = 1; j <= i; j++) {
        if (i % j == 0) {
          assert(s.count(j));
        }
      }
    }
  }
  {  // Compressed factors are convenient for building divisors.
    Factors factors{{2, 3}, {3, 2}, {5, 1}};
    vector<int64_t> divs = divisors_from_factors(factors);
    assert(
        (divs == vector<int64_t>{
                     1,  2,  3,  4,  5,  6,  8,  9,  10, 12,  15,  18,
                     20, 24, 30, 36, 40, 45, 60, 72, 90, 120, 180, 360,
                 })
    );
    assert(factorize(360) == factors);
    assert(divisors(360) == divs);
  }
  {  // Large factorization tests.
    const vector<int64_t> nums{
        (1LL << 62),                    // high power of two
        9223372036854775807LL,          // 2^63 - 1, many small factors
        999983LL * 999983,              // square of a large prime
        1900009LL * 1910009 * 2541547,  // three medium-size prime factors
        50000017LL * 50001037,          // 16-digit semiprime
    };
    cout << "Factorization test:" << endl;
    cout << setw(24) << left << "num:";
    cout << setw(20) << left << "factorize_slow():";
    cout << "factorize():" << endl;
    cout << fixed << setprecision(3);
    for (int64_t n : nums) {
      clock_t start = clock();
      Factors factors1 = factorize_slow(n);
      double t1 = static_cast<double>(clock() - start) / CLOCKS_PER_SEC;
      start = clock();
      Factors factors2 = factorize(n);
      double t2 = static_cast<double>(clock() - start) / CLOCKS_PER_SEC;
      validate(n, factors1);
      validate(n, factors2);
      assert(factors1 == factors2);
      cout << setw(26) << left << n << right << setw(8) << 1000 * t1 << "ms";
      cout << setw(16) << 1000 * t2 << "ms" << endl;
      cout << left;
    }
  }
  return 0;
}

Example Output

Primality test:
num:           is_prime_slow():  is_probable_prime():  is_prime():
772023803            0.006ms           0.005ms           0.001ms
2147483647           0.009ms           0.006ms           0.002ms
5705234089           0.014ms           0.000ms           0.001ms
6339503641           0.013ms           0.001ms           0.001ms
999966000289         0.171ms           0.001ms           0.001ms

Factorization test:
num:                    factorize_slow():   factorize():
4611686018427387904          0.002ms           0.002ms
9223372036854775807          0.049ms           0.007ms
999966000289                 0.503ms           0.043ms
9223361212852495307          0.952ms           0.096ms
2500052700017629            24.587ms           0.194ms
Back to top