:warning: library/math/ModularSqrt.hpp

Depends on

Required by

Code

#pragma once
#include "library/mod/Modint.hpp"
#include <optional>
#include <random>
#include <chrono>

template <typename T, T MOD>
bool is_quadratic_residue(Mint<T, MOD> a) {
    if (a == 0) return true;
    return a.pow((MOD - 1) / 2) == 1;
}

template <typename T, T MOD>
std::optional<Mint<T, MOD>> mod_sqrt(Mint<T, MOD> a) {
    if (a == 0) return Mint<T, MOD>(0);
    if (MOD == 2) return a;
    if (!is_quadratic_residue(a)) return std::nullopt;

    if (MOD % 4 == 3) {
        return a.pow((MOD + 1) / 4);
    }

    // Tonelli-Shanks
    long long s = 0, q = MOD - 1;
    while (q % 2 == 0) {
        q /= 2;
        s++;
    }

    // Find a non-quadratic residue z
    std::mt19937_64 rng(std::chrono::steady_clock::now().time_since_epoch().count());
    Mint<T, MOD> z;
    do {
        z = rng() % MOD;
    } while (is_quadratic_residue(z));

    long long m = s;
    Mint<T, MOD> c = z.pow(q);
    Mint<T, MOD> t = a.pow(q);
    Mint<T, MOD> r = a.pow((q + 1) / 2);

    while (t != 1) {
        if (t == 0) return Mint<T, MOD>(0);
        long long i = 0;
        Mint<T, MOD> temp = t;
        while (temp != 1) {
            temp *= temp;
            i++;
            if (i == m) return std::nullopt; // Should not happen for quadratic residues
        }

        Mint<T, MOD> b = c.pow(1LL << (m - i - 1));
        m = i;
        c = b * b;
        t *= c;
        r *= b;
    }
    return r;
}
#line 2 "library/math/ExtraGCD.hpp"
using ll = long long;
std::pair<ll, ll> ext_gcd(ll a, ll b) {
    if (b == 0)
        return {1, 0};
    auto [X, Y] = ext_gcd(b, a % b);
    // bX + (a%b)Y = gcd(a,b)
    // a%b = a - b(a/b)
    // ∴ aY + b(X-(a/b)Y) = gcd(a,b)
    ll x = Y, y = X - (a / b) * Y;
    return {x, y};
}
#line 3 "library/mod/Modint.hpp"
template <typename T, T MOD = 998244353> struct Mint {
    inline static constexpr T mod = MOD;
    T v;
    Mint() : v(0) {}
    Mint(signed v) : v(v) {}
    Mint(long long t) {
        v = t % MOD;
        if (v < 0)
            v += MOD;
    }

    static Mint raw(int v) {
        Mint x;
        x.v = v;
        return x;
    }

    Mint pow(long long k) const {
        Mint res(1), tmp(v);
        while (k) {
            if (k & 1)
                res *= tmp;
            tmp *= tmp;
            k >>= 1;
        }
        return res;
    }

    static Mint add_identity() { return Mint(0); }
    static Mint mul_identity() { return Mint(1); }

    // Mint inv()const{return pow(MOD-2);}
    Mint inv() const { return Mint(ext_gcd(v, mod).first); }

    Mint &operator+=(Mint a) {
        v += a.v;
        if (v >= MOD)
            v -= MOD;
        return *this;
    }
    Mint &operator-=(Mint a) {
        v += MOD - a.v;
        if (v >= MOD)
            v -= MOD;
        return *this;
    }
    Mint &operator*=(Mint a) {
        v = 1LL * v * a.v % MOD;
        return *this;
    }
    Mint &operator/=(Mint a) { return (*this) *= a.inv(); }

    Mint operator+(Mint a) const { return Mint(v) += a; }
    Mint operator-(Mint a) const { return Mint(v) -= a; }
    Mint operator*(Mint a) const { return Mint(v) *= a; }
    Mint operator/(Mint a) const { return Mint(v) /= a; }
#define FRIEND(op)                                                             \
    friend Mint operator op(int a, Mint b) { return Mint(a) op b; }
    FRIEND(+);
    FRIEND(-);
    FRIEND(*);
    FRIEND(/);
#undef FRIEND
    Mint operator+() const { return *this; }
    Mint operator-() const { return v ? Mint(MOD - v) : Mint(v); }

    bool operator==(const Mint a) const { return v == a.v; }
    bool operator!=(const Mint a) const { return v != a.v; }

    static Mint comb(long long n, int k) {
        Mint num(1), dom(1);
        for (int i = 0; i < k; i++) {
            num *= Mint(n - i);
            dom *= Mint(i + 1);
        }
        return num / dom;
    }

    friend std::ostream &operator<<(std::ostream &os, const Mint &m) {
        os << m.v;
        return os;
    }
    friend std::istream &operator>>(std::istream &is, Mint &m) {
        is >> m.v;
        m.v %= MOD;
        if (m.v < 0)
            m.v += MOD;
        return is;
    }
};
#line 3 "library/math/ModularSqrt.hpp"
#include <optional>
#include <random>
#include <chrono>

template <typename T, T MOD>
bool is_quadratic_residue(Mint<T, MOD> a) {
    if (a == 0) return true;
    return a.pow((MOD - 1) / 2) == 1;
}

template <typename T, T MOD>
std::optional<Mint<T, MOD>> mod_sqrt(Mint<T, MOD> a) {
    if (a == 0) return Mint<T, MOD>(0);
    if (MOD == 2) return a;
    if (!is_quadratic_residue(a)) return std::nullopt;

    if (MOD % 4 == 3) {
        return a.pow((MOD + 1) / 4);
    }

    // Tonelli-Shanks
    long long s = 0, q = MOD - 1;
    while (q % 2 == 0) {
        q /= 2;
        s++;
    }

    // Find a non-quadratic residue z
    std::mt19937_64 rng(std::chrono::steady_clock::now().time_since_epoch().count());
    Mint<T, MOD> z;
    do {
        z = rng() % MOD;
    } while (is_quadratic_residue(z));

    long long m = s;
    Mint<T, MOD> c = z.pow(q);
    Mint<T, MOD> t = a.pow(q);
    Mint<T, MOD> r = a.pow((q + 1) / 2);

    while (t != 1) {
        if (t == 0) return Mint<T, MOD>(0);
        long long i = 0;
        Mint<T, MOD> temp = t;
        while (temp != 1) {
            temp *= temp;
            i++;
            if (i == m) return std::nullopt; // Should not happen for quadratic residues
        }

        Mint<T, MOD> b = c.pow(1LL << (m - i - 1));
        m = i;
        c = b * b;
        t *= c;
        r *= b;
    }
    return r;
}
Back to top page