math/prime_64bit.hpp

• View this file on GitHub
• Last update: 2024-07-03 00:53:53+09:00
• Include: #include "math/prime_64bit.hpp"

Depends on

• math/factorize.hpp

Verified with

• test/math/PrimitiveRoot.test.cpp

Code

#pragma once

#include <type_traits>
#include <utility>
#include <algorithm>
#include "factorize.hpp"

namespace noya2::internal64bit {

template<typename T>
constexpr T safe_mod(T a, T p){
    a %= p;
    if constexpr (std::is_signed_v<T> || std::is_same_v<T, __int128_t>){
        if (a < 0) a += p;
    }
    return a;
}

template<typename T, typename U>
constexpr T pow_mod_constexpr(T x, U n, T p){
    if (p == 1) return 0;
    x = safe_mod(x, p);
    T ret = 1;
    while (n != 0){
        if (n % 2 == 1){
            ret = U(ret) * x % p;
        }
        x = U(x) * x % p;
        n /= 2;
    }
    return ret;
}

// return {g, y}
//   g = gcd(x, p), y * x == 1 (mod p/g)
template<typename T>
constexpr std::pair<T, T> inv_gcd(T x, T p){
    x = safe_mod(x, p);
    if (x == 0) return {p, 0};
    T s = p, t = x;
    T m0 = 0, m1 = 1;
    while (t != 0){
        T q = s / t;
        s -= t * q;
        m0 -= m1 * q;
        std::swap(s, t);
        std::swap(m0, m1);
    }
    if (m0 < 0) m0 += p / s;
    return {s, m0};
}

// p must be prime
long long primitive_root_ll(long long p){
    if (p == 2) return 1;
    auto fs = fast_factorize::factorize(p - 1);
    fs.erase(std::unique(fs.begin(), fs.end()), fs.end());
    for (long long g = 2; ; g++){
        bool ok = true;
        for (auto &f : fs){
            if (pow_mod_constexpr<long long, __int128_t>(g, (p - 1) / f, p) == 1){
                ok = false;
                break;
            }
        }
        if (ok) return g;
    }
    exit(1);
}

} // namespace noya2

#line 2 "math/prime_64bit.hpp"

#include <type_traits>
#include <utility>
#include <algorithm>
#line 2 "math/factorize.hpp"

#include <stddef.h>
#include <stdint.h>
#line 6 "math/factorize.hpp"
#include <initializer_list>
#include <iostream>
#include <vector>
#line 10 "math/factorize.hpp"

namespace fast_factorize {

/*
    See : https://judge.yosupo.jp/submission/189742
*/

// ---- gcd ----

uint64_t gcd_stein_impl( uint64_t x, uint64_t y ) {
    if( x == y ) { return x; }
    const uint64_t a = y - x;
    const uint64_t b = x - y;
    const int n = __builtin_ctzll( b );
    const uint64_t s = x < y ? a : b;
    const uint64_t t = x < y ? x : y;
    return gcd_stein_impl( s >> n, t );
}

uint64_t gcd_stein( uint64_t x, uint64_t y ) {
    if( x == 0 ) { return y; }
    if( y == 0 ) { return x; }
    const int n = __builtin_ctzll( x );
    const int m = __builtin_ctzll( y );
    return gcd_stein_impl( x >> n, y >> m ) << ( n < m ? n : m );
}

// ---- is_prime ----

uint64_t mod_pow( uint64_t x, uint64_t y, uint64_t mod ) {
    uint64_t ret = 1;
    uint64_t acc = x;
    for( ; y; y >>= 1 ) {
        if( y & 1 ) {
            ret = __uint128_t(ret) * acc % mod;
        }
        acc = __uint128_t(acc) * acc % mod;
    }
    return ret;
}

bool miller_rabin( uint64_t n, const std::initializer_list<uint64_t>& as ) {
    return std::all_of( as.begin(), as.end(), [n]( uint64_t a ) {
        if( n <= a ) { return true; }

        int e = __builtin_ctzll( n - 1 );
        uint64_t z = mod_pow( a, ( n - 1 ) >> e, n );
        if( z == 1 || z == n - 1 ) { return true; }

        while( --e ) {
            z = __uint128_t(z) * z % n;
            if( z == 1 ) { return false; }
            if( z == n - 1 ) { return true; }
        }

        return false;
    });
}

bool is_prime( uint64_t n ) {
    if( n == 2 ) { return true; }
    if( n % 2 == 0 ) { return false; }
    if( n < 4759123141 ) { return miller_rabin( n, { 2, 7, 61 } ); }
    return miller_rabin( n, { 2, 325, 9375, 28178, 450775, 9780504, 1795265022 } );
}

// ---- Montgomery ----

class Montgomery {
    uint64_t mod;
    uint64_t R;
public:
    Montgomery( uint64_t n ) : mod(n), R(n) {
       for( size_t i = 0; i < 5; ++i ) {
          R *= 2 - mod * R;
       }
    }

    uint64_t fma( uint64_t a, uint64_t b, uint64_t c ) const {
        const __uint128_t d = __uint128_t(a) * b;
        const uint64_t    e = c + mod + ( d >> 64 );
        const uint64_t    f = uint64_t(d) * R;
        const uint64_t    g = ( __uint128_t(f) * mod ) >> 64;
        return e - g;
    }

    uint64_t mul( uint64_t a, uint64_t b ) const {
        return fma( a, b, 0 );
    }
};

// ---- Pollard's rho algorithm ----

uint64_t pollard_rho( uint64_t n ) {
    if( n % 2 == 0 ) { return 2; }
    const Montgomery m( n );

    constexpr uint64_t C1 = 1;
    constexpr uint64_t C2 = 2;
    constexpr uint64_t M = 512;

    uint64_t Z1 = 1;
    uint64_t Z2 = 2;
retry:
    uint64_t z1 = Z1;
    uint64_t z2 = Z2;
    for( size_t k = M; ; k *= 2 ) {
        const uint64_t x1 = z1 + n;
        const uint64_t x2 = z2 + n;
        for( size_t j = 0; j < k; j += M ) {
            const uint64_t y1 = z1;
            const uint64_t y2 = z2;

            uint64_t q1 = 1;
            uint64_t q2 = 2;
            z1 = m.fma( z1, z1, C1 );
            z2 = m.fma( z2, z2, C2 );
            for( size_t i = 0; i < M; ++i ) {
                const uint64_t t1 = x1 - z1;
                const uint64_t t2 = x2 - z2;
                z1 = m.fma( z1, z1, C1 );
                z2 = m.fma( z2, z2, C2 );
                q1 = m.mul( q1, t1 );
                q2 = m.mul( q2, t2 );
            }
            q1 = m.mul( q1, x1 - z1 );
            q2 = m.mul( q2, x2 - z2 );

            const uint64_t q3 = m.mul( q1, q2 );
            const uint64_t g3 = gcd_stein( n, q3 );
            if( g3 == 1 ) { continue; }
            if( g3 != n ) { return g3; }

            const uint64_t g1 = gcd_stein( n, q1 );
            const uint64_t g2 = gcd_stein( n, q2 );

            const uint64_t C = g1 != 1 ? C1 : C2;
            const uint64_t x = g1 != 1 ? x1 : x2;
            uint64_t       z = g1 != 1 ? y1 : y2;
            uint64_t       g = g1 != 1 ? g1 : g2;

            if( g == n ) {
                do {
                    z = m.fma( z, z, C );
                    g = gcd_stein( n, x - z );
                } while( g == 1 );
            }
            if( g != n ) {
                return g;
            }

            Z1 += 2;
            Z2 += 2;
            goto retry;
        }
    }
}

void factorize_impl( uint64_t n, std::vector<uint64_t>& ret ) {
    if( n <= 1 ) { return; }
    if( is_prime( n ) ) { ret.push_back( n ); return; }

    const uint64_t p = pollard_rho( n );

    factorize_impl( p, ret );
    factorize_impl( n / p, ret );
}

std::vector<uint64_t> factorize( uint64_t n ) {
    std::vector<uint64_t> ret;
    factorize_impl( n, ret );
    std::sort( ret.begin(), ret.end() );
    return ret;
}

} // namespace fast_factorize

namespace noya2 {

std::vector<std::pair<long long, int>> factorize(long long n){
    std::vector<std::pair<long long, int>> ans;
    auto ps = fast_factorize::factorize(n);
    int sz = ps.size();
    for (int l = 0, r = 0; l < sz; l = r){
        while (r < sz && ps[l] == ps[r]) r++;
        ans.emplace_back(ps[l], r-l);
    }
    return ans;
}

std::vector<long long> divisors(long long n){
    auto ps = fast_factorize::factorize(n);
    int sz = ps.size();
    std::vector<long long> ans = {1};
    for (int l = 0, r = 0; l < sz; l = r){
        while (r < sz && ps[l] == ps[r]) r++;
        int e = r - l;
        int len = ans.size();
        ans.reserve(len*(e+1));
        long long mul = ps[l];
        while (true){
            for (int i = 0; i < len; i++){
                ans.emplace_back(ans[i]*mul);
            }
            if (--e == 0) break;
            mul *= ps[l];
        }
    }
    return ans;
}

std::vector<long long> divisors(const std::vector<std::pair<long long, int>> &pes){
    std::vector<long long> ans = {1};
    for (auto [p, e] : pes){
        int len = ans.size();
        ans.reserve(len*(e+1));
        long long mul = p;
        while (true){
            for (int i = 0; i < len; i++){
                ans.emplace_back(ans[i]*mul);
            }
            if (--e == 0) break;
            mul *= p;
        }
    }
    return ans;
}

} // namespace noya2
#line 7 "math/prime_64bit.hpp"

namespace noya2::internal64bit {

template<typename T>
constexpr T safe_mod(T a, T p){
    a %= p;
    if constexpr (std::is_signed_v<T> || std::is_same_v<T, __int128_t>){
        if (a < 0) a += p;
    }
    return a;
}

template<typename T, typename U>
constexpr T pow_mod_constexpr(T x, U n, T p){
    if (p == 1) return 0;
    x = safe_mod(x, p);
    T ret = 1;
    while (n != 0){
        if (n % 2 == 1){
            ret = U(ret) * x % p;
        }
        x = U(x) * x % p;
        n /= 2;
    }
    return ret;
}

// return {g, y}
//   g = gcd(x, p), y * x == 1 (mod p/g)
template<typename T>
constexpr std::pair<T, T> inv_gcd(T x, T p){
    x = safe_mod(x, p);
    if (x == 0) return {p, 0};
    T s = p, t = x;
    T m0 = 0, m1 = 1;
    while (t != 0){
        T q = s / t;
        s -= t * q;
        m0 -= m1 * q;
        std::swap(s, t);
        std::swap(m0, m1);
    }
    if (m0 < 0) m0 += p / s;
    return {s, m0};
}

// p must be prime
long long primitive_root_ll(long long p){
    if (p == 2) return 1;
    auto fs = fast_factorize::factorize(p - 1);
    fs.erase(std::unique(fs.begin(), fs.end()), fs.end());
    for (long long g = 2; ; g++){
        bool ok = true;
        for (auto &f : fs){
            if (pow_mod_constexpr<long long, __int128_t>(g, (p - 1) / f, p) == 1){
                ok = false;
                break;
            }
        }
        if (ok) return g;
    }
    exit(1);
}

} // namespace noya2

Back to top page