sqrt_mod.hpp

• View this file on GitHub
• Last update: 2025-01-19 15:28:01+08:00
• Include: #include "sqrt_mod.hpp"

Depends on

• rng.hpp

Required by

• fps_sqrt.hpp

Verified with

• test/formal_power_series/sqrt_of_formal_power_series.0.test.cpp

Code

#pragma once

#include "rng.hpp"
#include <random>
#include <type_traits>
#include <vector>

template<typename Tp> inline std::vector<Tp> sqrt_mod_prime(Tp a) {
    // Bostan--Mori's algorithm
    const auto p = Tp::mod();
    if (p == 2 || a == 0) return {a};
    if (a.pow(p / 2) == -1) return {};
    if ((p & 3) == 3) {
        const auto b = a.pow((p + 1) / 4);
        return {b, -b};
    }
    xoshiro256starstar rng(std::random_device{}());
    std::uniform_int_distribution<std::decay_t<decltype(p)>> dis(2, p - 1);
    Tp t;
    do { t = dis(rng); } while ((t * t - a * 4).pow(p / 2) != -1);
    Tp k0 = 1, k1, k2 = -t, k3 = a;
    for (auto e = (p + 1) >> 1;;) {
        if (e & 1) {
            k0 = k1 - k0 * k2, k1 *= k3;
        } else {
            k1 = k0 * k3 - k1 * k2;
        }
        if ((e >>= 1) == 0) return {k0, -k0};
        k2 = k3 + k3 - k2 * k2, k3 *= k3;
    }
}

#line 2 "sqrt_mod.hpp"

#line 2 "rng.hpp"

#include <cstdint>
#include <limits>

// see: https://prng.di.unimi.it/xoshiro256starstar.c
// original license CC0 1.0
class xoshiro256starstar {
    using u64 = std::uint64_t;

    static inline u64 rotl(const u64 x, int k) { return (x << k) | (x >> (64 - k)); }

    u64 s_[4];

    u64 next() {
        const u64 res = rotl(s_[1] * 5, 7) * 9;
        const u64 t   = s_[1] << 17;
        s_[2] ^= s_[0];
        s_[3] ^= s_[1];
        s_[1] ^= s_[2];
        s_[0] ^= s_[3];
        s_[2] ^= t;
        s_[3] = rotl(s_[3], 45);
        return res;
    }

public:
    // see: https://prng.di.unimi.it/splitmix64.c
    // original license CC0 1.0
    explicit xoshiro256starstar(u64 seed) {
        for (int i = 0; i < 4; ++i) {
            u64 z = (seed += 0x9e3779b97f4a7c15);
            z     = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9;
            z     = (z ^ (z >> 27)) * 0x94d049bb133111eb;
            s_[i] = z ^ (z >> 31);
        }
    }
    // see: https://en.cppreference.com/w/cpp/named_req/UniformRandomBitGenerator
    using result_type = u64;
    static constexpr u64 min() { return std::numeric_limits<u64>::min(); }
    static constexpr u64 max() { return std::numeric_limits<u64>::max(); }
    u64 operator()() { return next(); }
};
#line 4 "sqrt_mod.hpp"
#include <random>
#include <type_traits>
#include <vector>

template<typename Tp> inline std::vector<Tp> sqrt_mod_prime(Tp a) {
    // Bostan--Mori's algorithm
    const auto p = Tp::mod();
    if (p == 2 || a == 0) return {a};
    if (a.pow(p / 2) == -1) return {};
    if ((p & 3) == 3) {
        const auto b = a.pow((p + 1) / 4);
        return {b, -b};
    }
    xoshiro256starstar rng(std::random_device{}());
    std::uniform_int_distribution<std::decay_t<decltype(p)>> dis(2, p - 1);
    Tp t;
    do { t = dis(rng); } while ((t * t - a * 4).pow(p / 2) != -1);
    Tp k0 = 1, k1, k2 = -t, k3 = a;
    for (auto e = (p + 1) >> 1;;) {
        if (e & 1) {
            k0 = k1 - k0 * k2, k1 *= k3;
        } else {
            k1 = k0 * k3 - k1 * k2;
        }
        if ((e >>= 1) == 0) return {k0, -k0};
        k2 = k3 + k3 - k2 * k2, k3 *= k3;
    }
}

Back to top page