test/Math/primality_test/yosupo-primality_test.cpp

View this file on GitHub
Last update: 2025-04-13 19:22:19+09:00
Problem: https://judge.yosupo.jp/problem/primality_test

Depends on

Code

// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/primality_test

#include "Math/primality_test.hpp"
#include <bits/stdc++.h>

using namespace std;

int main()
{
    long long Q;
    cin >> Q;

     for (int q = 0; q < Q; q++)
     {
         long long N;
         cin >> N;
         if (PrimalityTest::is_prime(N))
         {
             cout << "Yes" << endl;
         }
         else
         {
             cout << "No" << endl;
         }
     }

    return 0;
}

#line 1 "test/Math/primality_test/yosupo-primality_test.cpp"
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/primality_test

#line 1 "Math/primality_test.hpp"
#include <bits/stdc++.h>
using namespace std;

#line 2 "Math/mod_power.hpp"
using namespace std;

#line 2 "Other/fast_power.hpp"
using namespace std;

template <class S>
S fast_pow(S x, long long n, function<S(S, S)> mul, function<S()> e)
{
    assert(0 <= n);

    S ans = e();

    while (n)
    {
        if (n & 1)
        {
            ans = mul(ans, x);
        }
        x = mul(x, x);
        n >>= 1;
    }

    return ans;
}
#line 5 "Math/mod_power.hpp"

template <typename T>
enable_if_t<is_integral_v<T> || is_same_v<T, __int128_t>, T>
mod_pow(T x, T n, T mod)
{
    assert(0 <= n);
    assert(0 < mod);
    assert(mod <= numeric_limits<T>::max() / mod);

    x %= mod;
    if (x < 0)
    {
        x += mod;
    }

    auto mul = [&](T a, T b) -> T
    {
        return (a * b) % mod;
    };

    auto e = [&]() -> T
    {
        return 1;
    };

    return fast_pow<T>(x, n, mul, e);
}
#line 5 "Math/primality_test.hpp"

struct PrimalityTest
{
private:
    static bool miller_rabin(long long N, const vector<long long> &A)
    {
        const int s = countr_zero((unsigned long long)N - 1);
        const long long d = (N - 1) >> s;

        auto is_composite = [&](long long a) -> bool
        {
            // Fermat の小定理が成り立たない
            if (a % N == 0)
            {
                return false;
            }

            long long x = mod_pow<__int128_t>(a, d, N);
            if (x == 1 || x == N - 1)
            {
                return false;
            }

            for (int t = 0; t < s - 1; ++t)
            {
                x = __int128_t(x) * x % N;
                if (x == N - 1)
                {
                    return false;
                }
            }

            return true;
        };

        return none_of(A.begin(), A.end(), is_composite);
    }

public:
    static bool is_prime(long long N)
    {
        assert(0 <= N);
        if (N <= 2)
        {
            return N == 2;
        }
        if (N % 2 == 0)
        {
            return false;
        }

        return miller_rabin(N, {2, 325, 9375, 28178, 450775, 9780504, 1795265022});
    }
};
#line 5 "test/Math/primality_test/yosupo-primality_test.cpp"

using namespace std;

int main()
{
    long long Q;
    cin >> Q;

     for (int q = 0; q < Q; q++)
     {
         long long N;
         cin >> N;
         if (PrimalityTest::is_prime(N))
         {
             cout << "Yes" << endl;
         }
         else
         {
             cout << "No" << endl;
         }
     }

    return 0;
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	all_prime_00	AC	780 ms	3 MB
g++	carmichael_00	AC	8 ms	3 MB
g++	example_00	AC	4 ms	3 MB
g++	hack_issue996_00	AC	4 ms	3 MB
g++	less_1000000000_00	AC	166 ms	3 MB
g++	prod_two_prime_00	AC	249 ms	3 MB
g++	random_00	AC	283 ms	3 MB
g++	random_01	AC	199 ms	3 MB
g++	random_02	AC	205 ms	3 MB
g++	small_00	AC	135 ms	3 MB