cp-algo/number_theory/discrete_sqrt.hpp

• View this file on GitHub
• Last update: 2024-10-07 21:21:13+02:00
• Include: #include "cp-algo/number_theory/discrete_sqrt.hpp"
• Link: https://en.wikipedia.org/wiki/Berlekamp-Rabin_algorithm

Depends on

• cp-algo/math/affine.hpp
• cp-algo/math/common.hpp
• cp-algo/number_theory/modint.hpp
• cp-algo/random/rng.hpp

Required by

• cp-algo/linalg/frobenius.hpp
• cp-algo/math/poly.hpp
• cp-algo/number_theory.hpp

Verified with

• Characteristic Polynomial (verify/linalg/characteristic.test.cpp)
• Pow of Matrix (Frobenius) (verify/linalg/pow_fast.test.cpp)
• Sqrt Mod (verify/number_theory/modsqrt.test.cpp)
• Division of Polynomials (verify/poly/div.test.cpp)
• Exp of Power Series (verify/poly/exp.test.cpp)
• Find Linear Recurrence (verify/poly/find_linrec.test.cpp)
• Polynomial Interpolation (verify/poly/inter.test.cpp)
• Inv of Power Series (verify/poly/inv.test.cpp)
• Inv of Polynomials (verify/poly/inv_mod.test.cpp)
• Consecutive Terms of Linear Recursion (verify/poly/linrec_consecutive.test.cpp)
• Kth term of Linearly Recurrent Sequence (verify/poly/linrec_single.test.cpp)
• Log of Power Series (verify/poly/log.test.cpp)
• Pow of Power Series (verify/poly/pow.test.cpp)
• Product of Polynomial Sequence (verify/poly/prod_sequence.test.cpp)
• Polynomial Root Finding (verify/poly/roots.test.cpp)
• Sqrt of Power Series (verify/poly/sqrt.test.cpp)
• Polynomial Taylor Shift (verify/poly/taylor.test.cpp)

Code

#ifndef CP_ALGO_NUMBER_THEORY_DISCRETE_SQRT_HPP
#define CP_ALGO_NUMBER_THEORY_DISCRETE_SQRT_HPP
#include "modint.hpp"
#include "../random/rng.hpp"
#include "../math/affine.hpp"
namespace cp_algo::math {
    // https://en.wikipedia.org/wiki/Berlekamp-Rabin_algorithm
    template<modint_type base>
    std::optional<base> sqrt(base b) {
        if(b == base(0)) {
            return base(0);
        } else if(bpow(b, (b.mod() - 1) / 2) != base(1)) {
            return std::nullopt;
        } else {
            while(true) {
                base z = random::rng();
                if(z * z == b) {
                    return z;
                }
                lin<base> x(1, z, b); // x + z (mod x^2 - b)
                x = bpow(x, (b.mod() - 1) / 2, lin<base>(0, 1, b));
                if(x.a != base(0)) {
                    return x.a.inv();
                }
            }
        }
    }
}
#endif // CP_ALGO_NUMBER_THEORY_SQRT_HPP

#line 1 "cp-algo/number_theory/discrete_sqrt.hpp"


#line 1 "cp-algo/number_theory/modint.hpp"


#line 1 "cp-algo/math/common.hpp"


#include <functional>
#include <cstdint>
namespace cp_algo::math {
#ifdef CP_ALGO_MAXN
    const int maxn = CP_ALGO_MAXN;
#else
    const int maxn = 1 << 19;
#endif
    const int magic = 64; // threshold for sizes to run the naive algo

    auto bpow(auto const& x, int64_t n, auto const& one, auto op) {
        if(n == 0) {
            return one;
        } else {
            auto t = bpow(x, n / 2, one, op);
            t = op(t, t);
            if(n % 2) {
                t = op(t, x);
            }
            return t;
        }
    }
    auto bpow(auto x, int64_t n, auto ans) {
        return bpow(x, n, ans, std::multiplies{});
    }
    template<typename T>
    T bpow(T const& x, int64_t n) {
        return bpow(x, n, T(1));
    }
}

#line 4 "cp-algo/number_theory/modint.hpp"
#include <iostream>
#include <cassert>
namespace cp_algo::math {
    inline constexpr uint64_t inv64(uint64_t x) {
        assert(x % 2);
        uint64_t y = 1;
        while(y * x != 1) {
            y *= 2 - x * y;
        }
        return y;
    }

    template<typename modint>
    struct modint_base {
        static int64_t mod() {
            return modint::mod();
        }
        static uint64_t imod() {
            return modint::imod();
        }
        static __uint128_t pw128() {
            return modint::pw128();
        }
        static uint64_t m_reduce(__uint128_t ab) {
            if(mod() % 2 == 0) [[unlikely]] {
                return ab % mod();
            } else {
                uint64_t m = ab * imod();
                return (ab + __uint128_t(m) * mod()) >> 64;
            }
        }
        static uint64_t m_transform(uint64_t a) {
            if(mod() % 2 == 0) [[unlikely]] {
                return a;
            } else {
                return m_reduce(a * pw128());
            }
        }
        modint_base(): r(0) {}
        modint_base(int64_t rr): r(rr % mod()) {
            r = std::min(r, r + mod());
            r = m_transform(r);
        }
        modint inv() const {
            return bpow(to_modint(), mod() - 2);
        }
        modint operator - () const {
            modint neg;
            neg.r = std::min(-r, 2 * mod() - r);
            return neg;
        }
        modint& operator /= (const modint &t) {
            return to_modint() *= t.inv();
        }
        modint& operator *= (const modint &t) {
            r = m_reduce(__uint128_t(r) * t.r);
            return to_modint();
        }
        modint& operator += (const modint &t) {
            r += t.r; r = std::min(r, r - 2 * mod());
            return to_modint();
        }
        modint& operator -= (const modint &t) {
            r -= t.r; r = std::min(r, r + 2 * mod());
            return to_modint();
        }
        modint operator + (const modint &t) const {return modint(to_modint()) += t;}
        modint operator - (const modint &t) const {return modint(to_modint()) -= t;}
        modint operator * (const modint &t) const {return modint(to_modint()) *= t;}
        modint operator / (const modint &t) const {return modint(to_modint()) /= t;}
        // Why <=> doesn't work?..
        auto operator == (const modint_base &t) const {return getr() == t.getr();}
        auto operator != (const modint_base &t) const {return getr() != t.getr();}
        auto operator <= (const modint_base &t) const {return getr() <= t.getr();}
        auto operator >= (const modint_base &t) const {return getr() >= t.getr();}
        auto operator < (const modint_base &t) const {return getr() < t.getr();}
        auto operator > (const modint_base &t) const {return getr() > t.getr();}
        int64_t rem() const {
            uint64_t R = getr();
            return 2 * R > (uint64_t)mod() ? R - mod() : R;
        }

        // Only use if you really know what you're doing!
        uint64_t modmod() const {return 8ULL * mod() * mod();};
        void add_unsafe(uint64_t t) {r += t;}
        void pseudonormalize() {r = std::min(r, r - modmod());}
        modint const& normalize() {
            if(r >= (uint64_t)mod()) {
                r %= mod();
            }
            return to_modint();
        }
        void setr(uint64_t rr) {r = m_transform(rr);}
        uint64_t getr() const {
            uint64_t res = m_reduce(r);
            return std::min(res, res - mod());
        }
        void setr_direct(uint64_t rr) {r = rr;}
        uint64_t getr_direct() const {return std::min(r, r - mod());}
    private:
        uint64_t r;
        modint& to_modint() {return static_cast<modint&>(*this);}
        modint const& to_modint() const {return static_cast<modint const&>(*this);}
    };
    template<typename modint>
    std::istream& operator >> (std::istream &in, modint_base<modint> &x) {
        uint64_t r;
        auto &res = in >> r;
        x.setr(r);
        return res;
    }
    template<typename modint>
    std::ostream& operator << (std::ostream &out, modint_base<modint> const& x) {
        return out << x.getr();
    }

    template<typename modint>
    concept modint_type = std::is_base_of_v<modint_base<modint>, modint>;

    template<int64_t m>
    struct modint: modint_base<modint<m>> {
        static constexpr uint64_t im = m % 2 ? inv64(-m) : 0;
        static constexpr uint64_t r2 = __uint128_t(-1) % m + 1;
        static constexpr int64_t mod() {return m;}
        static constexpr uint64_t imod() {return im;}
        static constexpr __uint128_t pw128() {return r2;}
        using Base = modint_base<modint<m>>;
        using Base::Base;
    };

    struct dynamic_modint: modint_base<dynamic_modint> {
        static int64_t mod() {return m;}
        static uint64_t imod() {return im;}
        static __uint128_t pw128() {return r2;}
        static void switch_mod(int64_t nm) {
            m = nm;
            im = m % 2 ? inv64(-m) : 0;
            r2 = __uint128_t(-1) % m + 1;
        }
        using Base = modint_base<dynamic_modint>;
        using Base::Base;

        // Wrapper for temp switching
        auto static with_mod(int64_t tmp, auto callback) {
            struct scoped {
                int64_t prev = mod();
                ~scoped() {switch_mod(prev);}
            } _;
            switch_mod(tmp);
            return callback();
        }
    private:
        static int64_t m;
        static uint64_t im, r1, r2;
    };
    int64_t dynamic_modint::m = 1;
    uint64_t dynamic_modint::im = -1;
    uint64_t dynamic_modint::r2 = 0;
}

#line 1 "cp-algo/random/rng.hpp"


#include <chrono>
#include <random>
namespace cp_algo::random {
    uint64_t rng() {
        static std::mt19937_64 rng(
            std::chrono::steady_clock::now().time_since_epoch().count()
        );
        return rng();
    }
}

#line 1 "cp-algo/math/affine.hpp"


#include <optional>
#include <utility>
#line 6 "cp-algo/math/affine.hpp"
#include <tuple>
namespace cp_algo::math {
    // a * x + b
    template<typename base>
    struct lin {
        base a = 1, b = 0;
        std::optional<base> c;
        lin() {}
        lin(base b): a(0), b(b) {}
        lin(base a, base b): a(a), b(b) {}
        lin(base a, base b, base _c): a(a), b(b), c(_c) {}

        // polynomial product modulo x^2 - c
        lin operator * (const lin& t) {
            assert(c && t.c && *c == *t.c);
            return {a * t.b + b * t.a, b * t.b + a * t.a * (*c), *c};
        }

        // a * (t.a * x + t.b) + b
        lin apply(lin const& t) const {
            return {a * t.a, a * t.b + b};
        }

        void prepend(lin const& t) {
            *this = t.apply(*this);
        }

        base eval(base x) const {
            return a * x + b;
        }
    };

    // (ax+b) / (cx+d)
    template<typename base>
    struct linfrac {
        base a, b, c, d;
        linfrac(): a(1), b(0), c(0), d(1) {} // x, identity for composition
        linfrac(base a): a(a), b(1), c(1), d(0) {} // a + 1/x, for continued fractions
        linfrac(base a, base b, base c, base d): a(a), b(b), c(c), d(d) {}

        // composition of two linfracs
        linfrac operator * (linfrac t) const {
            return t.prepend(linfrac(*this));
        }

        linfrac operator-() const {
            return {-a, -b, -c, -d};
        }

        linfrac adj() const {
            return {d, -b, -c, a};
        }
        
        linfrac& prepend(linfrac const& t) {
            t.apply(a, c);
            t.apply(b, d);
            return *this;
        }

        // apply linfrac to A/B
        void apply(base &A, base &B) const {
            std::tie(A, B) = std::pair{a * A + b * B, c * A + d * B};
        }
    };
}

#line 6 "cp-algo/number_theory/discrete_sqrt.hpp"
namespace cp_algo::math {
    // https://en.wikipedia.org/wiki/Berlekamp-Rabin_algorithm
    template<modint_type base>
    std::optional<base> sqrt(base b) {
        if(b == base(0)) {
            return base(0);
        } else if(bpow(b, (b.mod() - 1) / 2) != base(1)) {
            return std::nullopt;
        } else {
            while(true) {
                base z = random::rng();
                if(z * z == b) {
                    return z;
                }
                lin<base> x(1, z, b); // x + z (mod x^2 - b)
                x = bpow(x, (b.mod() - 1) / 2, lin<base>(0, 1, b));
                if(x.a != base(0)) {
                    return x.a.inv();
                }
            }
        }
    }
}

Back to top page