CP-Algorithms Library

This documentation is automatically generated by competitive-verifier/competitive-verifier

View the Project on GitHub cp-algorithms/cp-algorithms-aux

:warning: cp-algo/number_theory/two_squares.hpp

Depends on

Required by

Code

#ifndef CP_ALGO_NUMBER_THEORY_TWO_SQUARES_HPP
#define CP_ALGO_NUMBER_THEORY_TWO_SQUARES_HPP
#include "euler.hpp"
#include "../util/complex.hpp"
#include <cassert>
#include <utility>
#include <vector>
#include <map>
namespace cp_algo::math {
    using gaussint = complex<int64_t>;
    gaussint two_squares_prime_any(int64_t p) {
        if(p == 2) {
            return gaussint(1, 1);
        }
        assert(p % 4 == 1);
        using base = dynamic_modint<>;
        return base::with_mod(p, [&](){
            base g = primitive_root(p);
            int64_t i = bpow(g, (p - 1) / 4).getr();
            int64_t q0 = 1, q1 = 0;
            int64_t r = i, m = p;
            // TODO: Use library contfrac?
            do {
                int64_t d = r / m;
                q0 = std::exchange(q1, q0 + d * q1);
                r = std::exchange(m, r % m);
            } while(q1 < p / q1);
            return gaussint(q0, (base(i) * base(q0)).rem());
        });
    }

    std::vector<gaussint> two_squares_all(int64_t n) {
        if(n == 0) {
            return {0};
        }
        auto primes = factorize(n);
        std::map<int64_t, int> cnt;
        for(auto p: primes) {
            cnt[p]++;
        }
        std::vector<gaussint> res = {1};
        for(auto [p, c]: cnt) {
            std::vector<gaussint> nres;
            if(p % 4 == 3) {
                if(c % 2 == 0) {
                    auto mul = bpow(gaussint(p), c / 2);
                    for(auto p: res) {
                        nres.push_back(p * mul);
                    }
                }
            } else if(p % 4 == 1) {
                gaussint base = two_squares_prime_any(p);
                for(int i = 0; i <= c; i++) {
                    auto mul = bpow(base, i) * bpow(conj(base), c - i);
                    for(auto p: res) {
                        nres.push_back(p * mul);
                    }
                }
            } else if(p % 4 == 2) {
                auto mul = bpow(gaussint(1, 1), c);
                for(auto p: res) {
                    nres.push_back(p * mul);
                }
            }
            res = nres;
        }
        std::vector<gaussint> nres;
        for(auto p: res) {
            while(p.real() < 0 || p.imag() < 0) {
                p *= gaussint(0, 1);
            }
            nres.push_back(p);
            if(!p.real() || !p.imag()) {
                nres.emplace_back(p.imag(), p.real());
            }
        }
        return nres;
    }
}
#endif // CP_ALGO_NUMBER_THEORY_TWO_SQUARES_HPP

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


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


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


#line 1 "cp-algo/number_theory/primality.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, auto 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, auto n, auto ans) {
        return bpow(x, n, ans, std::multiplies{});
    }
    template<typename T>
    T bpow(T const& x, auto 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 auto inv2(auto x) {
        assert(x % 2);
        std::make_unsigned_t<decltype(x)> y = 1;
        while(y * x != 1) {
            y *= 2 - x * y;
        }
        return y;
    }

    template<typename modint, typename _Int>
    struct modint_base {
        using Int = _Int;
        using UInt = std::make_unsigned_t<Int>;
        static constexpr size_t bits = sizeof(Int) * 8;
        using Int2 = std::conditional_t<bits <= 32, int64_t, __int128_t>;
        using UInt2 = std::conditional_t<bits <= 32, uint64_t, __uint128_t>;
        static Int mod() {
            return modint::mod();
        }
        static UInt imod() {
            return modint::imod();
        }
        static UInt2 pw128() {
            return modint::pw128();
        }
        static UInt m_reduce(UInt2 ab) {
            if(mod() % 2 == 0) [[unlikely]] {
                return UInt(ab % mod());
            } else {
                UInt2 m = (UInt)ab * imod();
                return UInt((ab + m * mod()) >> bits);
            }
        }
        static UInt m_transform(UInt a) {
            if(mod() % 2 == 0) [[unlikely]] {
                return a;
            } else {
                return m_reduce(a * pw128());
            }
        }
        modint_base(): r(0) {}
        modint_base(Int2 rr): r(UInt(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((UInt2)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();}
        Int rem() const {
            UInt R = getr();
            return 2 * R > (UInt)mod() ? R - mod() : R;
        }

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

    template<auto m>
    struct modint: modint_base<modint<m>, decltype(m)> {
        using Base = modint_base<modint<m>, decltype(m)>;
        using Base::Base;
        static constexpr Base::UInt im = m % 2 ? inv2(-m) : 0;
        static constexpr Base::UInt r2 = (typename Base::UInt2)(-1) % m + 1;
        static constexpr Base::Int mod() {return m;}
        static constexpr Base::UInt imod() {return im;}
        static constexpr Base::UInt2 pw128() {return r2;}
    };

    template<typename Int = int64_t>
    struct dynamic_modint: modint_base<dynamic_modint<Int>, Int> {
        using Base = modint_base<dynamic_modint<Int>, Int>;
        using Base::Base;
        static Int mod() {return m;}
        static Base::UInt imod() {return im;}
        static Base::UInt2 pw128() {return r2;}
        static void switch_mod(Int nm) {
            m = nm;
            im = m % 2 ? inv2(-m) : 0;
            r2 = static_cast<Base::UInt>(static_cast<Base::UInt2>(-1) % m + 1);
        }

        // Wrapper for temp switching
        auto static with_mod(Int tmp, auto callback) {
            struct scoped {
                Int prev = mod();
                ~scoped() {switch_mod(prev);}
            } _;
            switch_mod(tmp);
            return callback();
        }
    private:
        static thread_local Int m;
        static thread_local Base::UInt im, r2;
    };
    template<typename Int>
    Int thread_local dynamic_modint<Int>::m = 1;
    template<typename Int>
    dynamic_modint<Int>::Base::UInt thread_local dynamic_modint<Int>::im = -1;
    template<typename Int>
    dynamic_modint<Int>::Base::UInt thread_local dynamic_modint<Int>::r2 = 0;
}

#line 4 "cp-algo/number_theory/primality.hpp"
#include <algorithm>
#include <bit>
namespace cp_algo::math {
    // https://en.wikipedia.org/wiki/Miller–Rabin_primality_test
    bool is_prime(uint64_t m) {
        if(m == 1 || m % 2 == 0) {
            return m == 2;
        }
        // m - 1 = 2^s * d
        int s = std::countr_zero(m - 1);
        auto d = (m - 1) >> s;
        using base = dynamic_modint<>;
        auto test = [&](base x) {
            x = bpow(x, d);
            if(std::abs(x.rem()) <= 1) {
                return true;
            }
            for(int i = 1; i < s && x != -1; i++) {
                x *= x;
            }
            return x == -1;
        };
        return base::with_mod(m, [&](){
            // Works for all m < 2^64: https://miller-rabin.appspot.com
            return std::ranges::all_of(std::array{
                2, 325, 9375, 28178, 450775, 9780504, 1795265022
            }, test);
        });
    }
}

#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 5 "cp-algo/number_theory/factorize.hpp"
#include <generator>
namespace cp_algo::math {
    // https://en.wikipedia.org/wiki/Pollard%27s_rho_algorithm
    auto proper_divisor(uint64_t m) {
        using base = dynamic_modint<>;
        return m % 2 == 0 ? 2 : base::with_mod(m, [&]() {
            base t = random::rng();
            auto f = [&](auto x) {
                return x * x + t;
            };
            base x, y;
            base g = 1;
            while(g == 1) {
                for(int i = 0; i < 64; i++) {
                    x = f(x);
                    y = f(f(y));
                    if(x == y) [[unlikely]] {
                        t = random::rng();
                        x = y = 0;
                    } else {
                        base t = g * (x - y);
                        g = t == 0 ? g : t;
                    }
                }
                g = std::gcd(g.getr(), m);
            }
            return g.getr();
        });
    }
    std::generator<uint64_t> factorize(uint64_t m) {
        if(is_prime(m)) {
            co_yield m;
        } else if(m > 1) {
            auto g = proper_divisor(m);
            co_yield std::ranges::elements_of(factorize(g));
            co_yield std::ranges::elements_of(factorize(m / g));
        }
    }
}

#line 4 "cp-algo/number_theory/euler.hpp"
namespace cp_algo::math {
    int64_t euler_phi(int64_t m) {
        auto primes = to<std::vector>(factorize(m));
        std::ranges::sort(primes);
        auto [from, to] = std::ranges::unique(primes);
        primes.erase(from, to);
        int64_t ans = m;
        for(auto it: primes) {
            ans -= ans / it;
        }
        return ans;
    }
    template<modint_type base>
    int64_t period(base x) {
        auto ans = euler_phi(base::mod());
        base x0 = bpow(x, ans);
        for(auto t: factorize(ans)) {
            while(ans % t == 0 && x0 * bpow(x, ans / t) == x0) {
                ans /= t;
            }
        }
        return ans;
    }
    int64_t primitive_root(int64_t p) {
        using base = dynamic_modint<>;
        return base::with_mod(p, [p](){
            base t = 1;
            while(period(t) != p - 1) {
                t = random::rng();
            }
            return t.getr();
        });
    }
}

#line 1 "cp-algo/util/complex.hpp"


#include <cmath>
namespace cp_algo {
    // Custom implementation, since std::complex is UB on non-floating types
    template<typename T>
    struct complex {
        using value_type = T;
        T x, y;
        constexpr complex() {}
        constexpr complex(T x): x(x), y(0) {}
        constexpr complex(T x, T y): x(x), y(y) {}
        complex& operator *= (T t) {x *= t; y *= t; return *this;}
        complex& operator /= (T t) {x /= t; y /= t; return *this;}
        complex operator * (T t) const {return complex(*this) *= t;}
        complex operator / (T t) const {return complex(*this) /= t;}
        complex& operator += (complex t) {x += t.x; y += t.y; return *this;}
        complex& operator -= (complex t) {x -= t.x; y -= t.y; return *this;}
        complex operator * (complex t) const {return {x * t.x - y * t.y, x * t.y + y * t.x};}
        complex operator / (complex t) const {return *this * t.conj() / t.norm();}
        complex operator + (complex t) const {return complex(*this) += t;}
        complex operator - (complex t) const {return complex(*this) -= t;}
        complex& operator *= (complex t) {return *this = *this * t;}
        complex& operator /= (complex t) {return *this = *this / t;}
        complex operator - () const {return {-x, -y};}
        complex conj() const {return {x, -y};}
        T norm() const {return x * x + y * y;}
        T abs() const {return std::sqrt(norm());}
        T real() const {return x;}
        T imag() const {return y;}
        T& real() {return x;}
        T& imag() {return y;}
        static complex polar(T r, T theta) {return {r * cos(theta), r * sin(theta)};}
        auto operator <=> (complex const& t) const = default;
    };
    template<typename T>
    complex<T> operator * (auto x, complex<T> y) {return y *= x;}
    template<typename T> complex<T> conj(complex<T> x) {return x.conj();}
    template<typename T> T norm(complex<T> x) {return x.norm();}
    template<typename T> T abs(complex<T> x) {return x.abs();}
    template<typename T> T& real(complex<T> &x) {return x.real();}
    template<typename T> T& imag(complex<T> &x) {return x.imag();}
    template<typename T> T real(complex<T> const& x) {return x.real();}
    template<typename T> T imag(complex<T> const& x) {return x.imag();}
    template<typename T> complex<T> polar(T r, T theta) {return complex<T>::polar(r, theta);}
}

#line 6 "cp-algo/number_theory/two_squares.hpp"
#include <utility>
#include <vector>
#include <map>
namespace cp_algo::math {
    using gaussint = complex<int64_t>;
    gaussint two_squares_prime_any(int64_t p) {
        if(p == 2) {
            return gaussint(1, 1);
        }
        assert(p % 4 == 1);
        using base = dynamic_modint<>;
        return base::with_mod(p, [&](){
            base g = primitive_root(p);
            int64_t i = bpow(g, (p - 1) / 4).getr();
            int64_t q0 = 1, q1 = 0;
            int64_t r = i, m = p;
            // TODO: Use library contfrac?
            do {
                int64_t d = r / m;
                q0 = std::exchange(q1, q0 + d * q1);
                r = std::exchange(m, r % m);
            } while(q1 < p / q1);
            return gaussint(q0, (base(i) * base(q0)).rem());
        });
    }

    std::vector<gaussint> two_squares_all(int64_t n) {
        if(n == 0) {
            return {0};
        }
        auto primes = factorize(n);
        std::map<int64_t, int> cnt;
        for(auto p: primes) {
            cnt[p]++;
        }
        std::vector<gaussint> res = {1};
        for(auto [p, c]: cnt) {
            std::vector<gaussint> nres;
            if(p % 4 == 3) {
                if(c % 2 == 0) {
                    auto mul = bpow(gaussint(p), c / 2);
                    for(auto p: res) {
                        nres.push_back(p * mul);
                    }
                }
            } else if(p % 4 == 1) {
                gaussint base = two_squares_prime_any(p);
                for(int i = 0; i <= c; i++) {
                    auto mul = bpow(base, i) * bpow(conj(base), c - i);
                    for(auto p: res) {
                        nres.push_back(p * mul);
                    }
                }
            } else if(p % 4 == 2) {
                auto mul = bpow(gaussint(1, 1), c);
                for(auto p: res) {
                    nres.push_back(p * mul);
                }
            }
            res = nres;
        }
        std::vector<gaussint> nres;
        for(auto p: res) {
            while(p.real() < 0 || p.imag() < 0) {
                p *= gaussint(0, 1);
            }
            nres.push_back(p);
            if(!p.real() || !p.imag()) {
                nres.emplace_back(p.imag(), p.real());
            }
        }
        return nres;
    }
}
Back to top page