This documentation is automatically generated by online-judge-tools/verification-helper

View the Project on GitHub maspypy/library

#include "mod/mod_sqrt.hpp"

#include "random/base.hpp" #include "mod/mod_pow.hpp" // p は素数. 解なしは -1. int mod_sqrt(int a, int p) { if (p == 2) return a; if (a == 0) return 0; int k = (p - 1) / 2; if (mod_pow(a, k, p) != 1) return -1; auto find = [&]() -> pi { while (1) { ll b = RNG(2, p); ll D = (b * b - a) % p; if (D == 0) return {b, D}; if (mod_pow(D, k, p) != 1) return {b, D}; } }; auto [b, D] = find(); if (D == 0) return b; ++k; // (b + sqrt(D))^k ll f0 = b, f1 = 1, g0 = 1, g1 = 0; while (k) { if (k & 1) { tie(g0, g1) = mp(f0 * g0 + D * f1 % p * g1, f1 * g0 + f0 * g1); g0 %= p, g1 %= p; } tie(f0, f1) = mp(f0 * f0 + D * f1 % p * f1, 2 * f0 * f1); f0 %= p, f1 %= p; k >>= 1; } if (g0 < 0) g0 += p; return g0; } // p は素数. 解なしは -1. ll mod_sqrt_64(ll a, ll p) { if (p == 2) return a; if (a == 0) return 0; ll k = (p - 1) / 2; if (mod_pow_64(a, k, p) != 1) return -1; auto find = [&]() -> pair<i128, i128> { while (1) { i128 b = RNG(2, p); i128 D = b * b - a; if (D == 0) return {b, D}; if (mod_pow_64(D, k, p) != 1) return {b, D}; } }; auto [b, D] = find(); if (D == 0) return b; ++k; // (b + sqrt(D))^k i128 f0 = b, f1 = 1, g0 = 1, g1 = 0; while (k) { if (k & 1) { tie(g0, g1) = mp(f0 * g0 + D * f1 % p * g1, f1 * g0 + f0 * g1); g0 %= p, g1 %= p; } tie(f0, f1) = mp(f0 * f0 + D * f1 % p * f1, 2 * f0 * f1); f0 %= p, f1 %= p; k >>= 1; } return g0; }

#line 2 "random/base.hpp" u64 RNG_64() { static uint64_t x_ = uint64_t(chrono::duration_cast<chrono::nanoseconds>(chrono::high_resolution_clock::now().time_since_epoch()).count()) * 10150724397891781847ULL; x_ ^= x_ << 7; return x_ ^= x_ >> 9; } u64 RNG(u64 lim) { return RNG_64() % lim; } ll RNG(ll l, ll r) { return l + RNG_64() % (r - l); } #line 2 "mod/mod_pow.hpp" #line 2 "mod/mongomery_modint.hpp" // odd mod. // x の代わりに rx を持つ template <int id, typename U1, typename U2> struct Mongomery_modint { using mint = Mongomery_modint; inline static U1 m, r, n2; static constexpr int W = numeric_limits<U1>::digits; static void set_mod(U1 mod) { assert(mod & 1 && mod <= U1(1) << (W - 2)); m = mod, n2 = -U2(m) % m, r = m; FOR(5) r *= 2 - m * r; r = -r; assert(r * m == U1(-1)); } static U1 reduce(U2 b) { return (b + U2(U1(b) * r) * m) >> W; } U1 x; Mongomery_modint() : x(0) {} Mongomery_modint(U1 x) : x(reduce(U2(x) * n2)){}; U1 val() const { U1 y = reduce(x); return y >= m ? y - m : y; } mint &operator+=(mint y) { x = ((x += y.x) >= m ? x - m : x); return *this; } mint &operator-=(mint y) { x -= (x >= y.x ? y.x : y.x - m); return *this; } mint &operator*=(mint y) { x = reduce(U2(x) * y.x); return *this; } mint operator+(mint y) const { return mint(*this) += y; } mint operator-(mint y) const { return mint(*this) -= y; } mint operator*(mint y) const { return mint(*this) *= y; } bool operator==(mint y) const { return (x >= m ? x - m : x) == (y.x >= m ? y.x - m : y.x); } bool operator!=(mint y) const { return not operator==(y); } mint pow(ll n) const { assert(n >= 0); mint y = 1, z = *this; for (; n; n >>= 1, z *= z) if (n & 1) y *= z; return y; } }; template <int id> using Mongomery_modint_32 = Mongomery_modint<id, u32, u64>; template <int id> using Mongomery_modint_64 = Mongomery_modint<id, u64, u128>; #line 2 "mod/barrett.hpp" // https://github.com/atcoder/ac-library/blob/master/atcoder/internal_math.hpp struct Barrett { u32 m; u64 im; explicit Barrett(u32 m = 1) : m(m), im(u64(-1) / m + 1) {} u32 umod() const { return m; } u32 modulo(u64 z) { if (m == 1) return 0; u64 x = (u64)(((unsigned __int128)(z)*im) >> 64); u64 y = x * m; return (z - y + (z < y ? m : 0)); } u64 floor(u64 z) { if (m == 1) return z; u64 x = (u64)(((unsigned __int128)(z)*im) >> 64); u64 y = x * m; return (z < y ? x - 1 : x); } pair<u64, u32> divmod(u64 z) { if (m == 1) return {z, 0}; u64 x = (u64)(((unsigned __int128)(z)*im) >> 64); u64 y = x * m; if (z < y) return {x - 1, z - y + m}; return {x, z - y}; } u32 mul(u32 a, u32 b) { return modulo(u64(a) * b); } }; struct Barrett_64 { u128 mod, mh, ml; explicit Barrett_64(u64 mod = 1) : mod(mod) { u128 m = u128(-1) / mod; if (m * mod + mod == u128(0)) ++m; mh = m >> 64; ml = m & u64(-1); } u64 umod() const { return mod; } u64 modulo(u128 x) { u128 z = (x & u64(-1)) * ml; z = (x & u64(-1)) * mh + (x >> 64) * ml + (z >> 64); z = (x >> 64) * mh + (z >> 64); x -= z * mod; return x < mod ? x : x - mod; } u64 mul(u64 a, u64 b) { return modulo(u128(a) * b); } }; #line 5 "mod/mod_pow.hpp" u32 mod_pow(int a, ll n, int mod) { assert(n >= 0); a = ((a %= mod) < 0 ? a + mod : a); if ((mod & 1) && (mod < (1 << 30))) { using mint = Mongomery_modint_32<202311021>; mint::set_mod(mod); return mint(a).pow(n).val(); } Barrett bt(mod); int r = 1; while (n) { if (n & 1) r = bt.mul(r, a); a = bt.mul(a, a), n >>= 1; } return r; } u64 mod_pow_64(ll a, ll n, u64 mod) { assert(n >= 0); a = ((a %= mod) < 0 ? a + mod : a); if ((mod & 1) && (mod < (u64(1) << 62))) { using mint = Mongomery_modint_64<202311021>; mint::set_mod(mod); return mint(a).pow(n).val(); } Barrett_64 bt(mod); ll r = 1; while (n) { if (n & 1) r = bt.mul(r, a); a = bt.mul(a, a), n >>= 1; } return r; } #line 3 "mod/mod_sqrt.hpp" // p は素数. 解なしは -1. int mod_sqrt(int a, int p) { if (p == 2) return a; if (a == 0) return 0; int k = (p - 1) / 2; if (mod_pow(a, k, p) != 1) return -1; auto find = [&]() -> pi { while (1) { ll b = RNG(2, p); ll D = (b * b - a) % p; if (D == 0) return {b, D}; if (mod_pow(D, k, p) != 1) return {b, D}; } }; auto [b, D] = find(); if (D == 0) return b; ++k; // (b + sqrt(D))^k ll f0 = b, f1 = 1, g0 = 1, g1 = 0; while (k) { if (k & 1) { tie(g0, g1) = mp(f0 * g0 + D * f1 % p * g1, f1 * g0 + f0 * g1); g0 %= p, g1 %= p; } tie(f0, f1) = mp(f0 * f0 + D * f1 % p * f1, 2 * f0 * f1); f0 %= p, f1 %= p; k >>= 1; } if (g0 < 0) g0 += p; return g0; } // p は素数. 解なしは -1. ll mod_sqrt_64(ll a, ll p) { if (p == 2) return a; if (a == 0) return 0; ll k = (p - 1) / 2; if (mod_pow_64(a, k, p) != 1) return -1; auto find = [&]() -> pair<i128, i128> { while (1) { i128 b = RNG(2, p); i128 D = b * b - a; if (D == 0) return {b, D}; if (mod_pow_64(D, k, p) != 1) return {b, D}; } }; auto [b, D] = find(); if (D == 0) return b; ++k; // (b + sqrt(D))^k i128 f0 = b, f1 = 1, g0 = 1, g1 = 0; while (k) { if (k & 1) { tie(g0, g1) = mp(f0 * g0 + D * f1 % p * g1, f1 * g0 + f0 * g1); g0 %= p, g1 %= p; } tie(f0, f1) = mp(f0 * f0 + D * f1 % p * f1, 2 * f0 * f1); f0 %= p, f1 %= p; k >>= 1; } return g0; }