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

View the Project on GitHub maspypy/library

#include "nt/three_square.hpp"

#include "nt/primetest.hpp" #include "nt/gaussian_integers.hpp" // https://math.stackexchange.com/questions/483101/rabin-and-shallit-algorithm // ERH のもと O(log^2N) ？ tuple<ll, ll, ll> three_square(ll N) { if (N == 0) return {0, 0, 0}; auto F = [&](ll n) -> tuple<ll, ll, ll> { if (N == 2) return {1, 1, 0}; if (N == 3) return {1, 1, 1}; if (N == 10) return {3, 1, 0}; if (N == 34) return {5, 3, 0}; if (N == 58) return {7, 3, 0}; if (N == 85) return {9, 2, 0}; if (N == 130) return {11, 3, 0}; if (N == 214) return {14, 3, 3}; if (N == 226) return {15, 1, 0}; if (N == 370) return {19, 3, 0}; if (N == 526) return {21, 9, 2}; if (N == 706) return {25, 9, 0}; if (N == 730) return {27, 1, 0}; if (N == 1414) return {33, 18, 1}; if (N == 1906) return {41, 15, 0}; if (N == 2986) return {45, 31, 0}; if (N == 9634) return {97, 15, 0}; ll x = sqrtl(N); if (N == x * x) return {x, 0, 0}; if (N % 4 != 1 && x % 2 == 0) --x; if (N % 4 == 1 && x % 2 == 1) --x; x += 2; while (1) { x -= 2; ll k = N - x * x; if (k < 0) break; if (k % 2 == 1 && primetest(k)) { auto [a, b] = solve_norm_equation_prime(k); a = abs(a), b = abs(b); return {a, b, x}; } if (k % 2 == 0 && primetest(k / 2)) { auto [a, b] = solve_norm_equation_prime(k / 2); tie(a, b) = mp(a + b, a - b); a = abs(a), b = abs(b); return {a, b, x}; } } assert(0); return {-1, -1, -1}; }; ll e = 0; while (N % 4 == 0) N /= 4, ++e; if (N % 8 == 7) return {-1, -1, -1}; auto [a, b, c] = F(N); return {a << e, b << e, c << e}; }

#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 3 "nt/primetest.hpp" bool primetest(const u64 x) { assert(x < u64(1) << 62); if (x == 2 or x == 3 or x == 5 or x == 7) return true; if (x % 2 == 0 or x % 3 == 0 or x % 5 == 0 or x % 7 == 0) return false; if (x < 121) return x > 1; const u64 d = (x - 1) >> lowbit(x - 1); using mint = Mongomery_modint_64<202311020>; mint::set_mod(x); const mint one(u64(1)), minus_one(x - 1); auto ok = [&](u64 a) -> bool { auto y = mint(a).pow(d); u64 t = d; while (y != one && y != minus_one && t != x - 1) y *= y, t <<= 1; if (y != minus_one && t % 2 == 0) return false; return true; }; if (x < (u64(1) << 32)) { for (u64 a: {2, 7, 61}) if (!ok(a)) return false; } else { for (u64 a: {2, 325, 9375, 28178, 450775, 9780504, 1795265022}) { if (!ok(a)) return false; } } return true; } #line 2 "nt/factor.hpp" #line 2 "random/base.hpp" u64 RNG_64() { static u64 x_ = u64(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 5 "nt/factor.hpp" template <typename mint> ll rho(ll n, ll c) { assert(n > 1); const mint cc(c); auto f = [&](mint x) { return x * x + cc; }; mint x = 1, y = 2, z = 1, q = 1; ll g = 1; const ll m = 1LL << (__lg(n) / 5); for (ll r = 1; g == 1; r <<= 1) { x = y; FOR(r) y = f(y); for (ll k = 0; k < r && g == 1; k += m) { z = y; FOR(min(m, r - k)) y = f(y), q *= x - y; g = gcd(q.val(), n); } } if (g == n) do { z = f(z); g = gcd((x - z).val(), n); } while (g == 1); return g; } ll find_prime_factor(ll n) { assert(n > 1); if (primetest(n)) return n; FOR(100) { ll m = 0; if (n < (1 << 30)) { using mint = Mongomery_modint_32<20231025>; mint::set_mod(n); m = rho<mint>(n, RNG(0, n)); } else { using mint = Mongomery_modint_64<20231025>; mint::set_mod(n); m = rho<mint>(n, RNG(0, n)); } if (primetest(m)) return m; n = m; } assert(0); return -1; } // ソートしてくれる vc<pair<ll, int>> factor(ll n) { assert(n >= 1); vc<pair<ll, int>> pf; FOR(p, 2, 100) { if (p * p > n) break; if (n % p == 0) { ll e = 0; do { n /= p, e += 1; } while (n % p == 0); pf.eb(p, e); } } while (n > 1) { ll p = find_prime_factor(n); ll e = 0; do { n /= p, e += 1; } while (n % p == 0); pf.eb(p, e); } sort(all(pf)); return pf; } vc<pair<ll, int>> factor_by_lpf(ll n, vc<int>& lpf) { vc<pair<ll, int>> res; while (n > 1) { int p = lpf[n]; int e = 0; while (n % p == 0) { n /= p; ++e; } res.eb(p, e); } return res; } #line 2 "mod/mod_pow.hpp" #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 "nt/gaussian_integers.hpp" template <typename T> struct Gaussian_Integer { T x, y; using G = Gaussian_Integer; Gaussian_Integer(T x = 0, T y = 0) : x(x), y(y) {} Gaussian_Integer(pair<T, T> p) : x(p.fi), y(p.se) {} T norm() const { return x * x + y * y; } G conjugate() const { return G(x, -y); } G &operator+=(const G &g) { x += g.x, y += g.y; return *this; } G &operator-=(const G &g) { x -= g.x, y -= g.y; return *this; } G &operator*=(const G &g) { tie(x, y) = mp(x * g.x - y * g.y, x * g.y + y * g.x); return *this; } G &operator/=(const G &g) { *this *= g.conjugate(); T n = g.norm(); x = floor(x + n / 2, n); y = floor(y + n / 2, n); return *this; } G &operator%=(const G &g) { auto q = G(*this) / g; q *= g; (*this) -= q; return *this; } G operator-() { return G(-x, -y); } G operator+(const G &g) { return G(*this) += g; } G operator-(const G &g) { return G(*this) -= g; } G operator*(const G &g) { return G(*this) *= g; } G operator/(const G &g) { return G(*this) /= g; } G operator%(const G &g) { return G(*this) %= g; } bool operator==(const G &g) { return (x == g.x && y == g.y); } static G gcd(G a, G b) { while (b.x != 0 || b.y != 0) { a %= b; swap(a, b); } return a; } G pow(ll n) const { assert(n >= 0); G ret(1), mul(*this); while (n > 0) { if (n & 1) ret *= mul; mul *= mul; n >>= 1; } return ret; } // (g,x,y) s.t ax+by=g static tuple<G, G, G> extgcd(G a, G b) { if (b.x != 0 || b.y != 0) { G q = a / b; auto [g, x, y] = extgcd(b, a - q * b); return {g, y, x - q * y}; } return {a, G{1, 0}, G{0, 0}}; } }; pair<ll, ll> solve_norm_equation_prime(ll p) { using G = Gaussian_Integer<i128>; assert(p == 2 || p % 4 == 1); if (p == 2) return {1, 1}; ll x = [&]() -> ll { ll x = 1; while (1) { ++x; ll pow_x = 1; if (p < (1 << 30)) { pow_x = mod_pow(x, (p - 1) / 4, p); if (pow_x * pow_x % p == p - 1) return pow_x; } else { pow_x = mod_pow_64(x, (p - 1) / 4, p); if (i128(pow_x) * pow_x % p == p - 1) return pow_x; } } return -1; }(); G a(p, 0), b(x, 1); a = G::gcd(a, b); assert(a.norm() == p); return {a.x, a.y}; } template <typename T> vc<Gaussian_Integer<T>> solve_norm_equation_factor(vc<pair<ll, int>> pfs) { using G = Gaussian_Integer<T>; vc<G> res; for (auto &&[p, e]: pfs) { if (p % 4 == 3 && e % 2 == 1) return {}; } res.eb(G(1, 0)); for (auto &&[p, e]: pfs) { if (p % 4 == 3) { T pp = 1; FOR(e / 2) pp *= p; for (auto &&g: res) { g.x *= pp; g.y *= pp; } continue; } auto [pix, piy] = solve_norm_equation_prime(p); G pi(pix, piy); vc<G> pows(e + 1); pows[0] = G(1, 0); FOR(i, e) pows[i + 1] = pows[i] * pi; if (p == 2) { for (auto &&g: res) g *= pows[e]; continue; } vc<G> pis(e + 1); FOR(j, e + 1) { pis[j] = pows[j] * (pows[e - j].conjugate()); } vc<G> new_res; new_res.reserve(len(res) * (e + 1)); for (auto &&g: res) { for (auto &&a: pis) { new_res.eb(g * a); } } swap(res, new_res); } for (auto &&g: res) { while (g.x <= 0 || g.y < 0) { g = G(-g.y, g.x); } } return res; } // i128 を使うと N <= 10^{18} もできる // ノルムがとれるように、2 乗してもオーバーフローしない型を使おう // 0 <= arg < 90 となるもののみ返す。 // 単数倍は作らないので、使うときに気を付ける。 template <typename T> vc<Gaussian_Integer<T>> solve_norm_equation(T N) { using G = Gaussian_Integer<T>; vc<G> res; if (N < 0) return {}; if (N == 0) { res.eb(G(0, 0)); return res; } auto pfs = factor(N); return solve_norm_equation_factor<T>(pfs); } #line 3 "nt/three_square.hpp" // https://math.stackexchange.com/questions/483101/rabin-and-shallit-algorithm // ERH のもと O(log^2N) ？ tuple<ll, ll, ll> three_square(ll N) { if (N == 0) return {0, 0, 0}; auto F = [&](ll n) -> tuple<ll, ll, ll> { if (N == 2) return {1, 1, 0}; if (N == 3) return {1, 1, 1}; if (N == 10) return {3, 1, 0}; if (N == 34) return {5, 3, 0}; if (N == 58) return {7, 3, 0}; if (N == 85) return {9, 2, 0}; if (N == 130) return {11, 3, 0}; if (N == 214) return {14, 3, 3}; if (N == 226) return {15, 1, 0}; if (N == 370) return {19, 3, 0}; if (N == 526) return {21, 9, 2}; if (N == 706) return {25, 9, 0}; if (N == 730) return {27, 1, 0}; if (N == 1414) return {33, 18, 1}; if (N == 1906) return {41, 15, 0}; if (N == 2986) return {45, 31, 0}; if (N == 9634) return {97, 15, 0}; ll x = sqrtl(N); if (N == x * x) return {x, 0, 0}; if (N % 4 != 1 && x % 2 == 0) --x; if (N % 4 == 1 && x % 2 == 1) --x; x += 2; while (1) { x -= 2; ll k = N - x * x; if (k < 0) break; if (k % 2 == 1 && primetest(k)) { auto [a, b] = solve_norm_equation_prime(k); a = abs(a), b = abs(b); return {a, b, x}; } if (k % 2 == 0 && primetest(k / 2)) { auto [a, b] = solve_norm_equation_prime(k / 2); tie(a, b) = mp(a + b, a - b); a = abs(a), b = abs(b); return {a, b, x}; } } assert(0); return {-1, -1, -1}; }; ll e = 0; while (N % 4 == 0) N /= 4, ++e; if (N % 8 == 7) return {-1, -1, -1}; auto [a, b, c] = F(N); return {a << e, b << e, c << e}; }