Formal power series (Sqrt)
(Mylib/Math/fps_sqrt.cpp)
Operations
Requirements
Notes
Problems
References
Depends on
Verified with
Code
#pragma once
#include "Mylib/Math/formal_power_series.cpp"
#include "Mylib/Number/Mod/mod_sqrt.cpp"
namespace haar_lib {
template <typename T, const auto &convolve>
auto formal_power_series<T, convolve>::sqrt() const -> std::optional<formal_power_series<T, convolve>> {
constexpr int mod = value_type::mod();
const int n = data_.size();
int k = 0;
for (; k < n; ++k)
if (data_[k] != 0) break;
if (k >= n) return *this;
if (k % 2 != 0) return {};
auto x = mod_sqrt((int64_t) data_[k], mod);
if (not x) return {};
const int m = n - k;
auto it = data_.begin() + k;
formal_power_series ret({*x});
int t = 1;
while (t <= m * 2) {
formal_power_series f(std::vector<T>(it, it + std::min(t, m)));
ret.resize(t);
f.resize(t);
ret = (ret + f * ret.inv()) * T(2).inv();
t <<= 1;
}
ret.resize(n);
ret = ret.shift(k / 2);
return ret;
}
} // namespace haar_lib
#line 2 "Mylib/Math/formal_power_series.cpp"
#include <cassert>
#include <functional>
#include <initializer_list>
#include <vector>
namespace haar_lib {
template <typename T, const auto &convolve>
class formal_power_series {
public:
using value_type = T;
private:
std::vector<T> data_;
public:
formal_power_series() {}
explicit formal_power_series(int N) : data_(N) {}
formal_power_series(const std::vector<T> &data_) : data_(data_) {}
formal_power_series(std::initializer_list<T> init) : data_(init.begin(), init.end()) {}
formal_power_series(const formal_power_series &a) : data_(a.data_) {}
formal_power_series(formal_power_series &&a) noexcept { *this = std::move(a); }
size_t size() const {
return data_.size();
}
const T &operator[](int i) const {
return data_[i];
}
T &operator[](int i) {
return data_[i];
}
auto begin() { return data_.begin(); }
auto end() { return data_.end(); }
const auto &data() const { return data_; }
void resize(int n) {
data_.resize(n);
}
auto &operator=(formal_power_series &&rhs) noexcept {
if (this != &rhs) {
data_ = std::move(rhs.data_);
}
return *this;
}
auto &operator+=(const formal_power_series &rhs) {
if (data_.size() < rhs.size()) data_.resize(rhs.size());
for (int i = 0; i < rhs.size(); ++i) data_[i] += rhs[i];
return *this;
}
auto &operator+=(T rhs) {
data_[0] += rhs;
return *this;
}
auto operator+(T rhs) const {
auto ret = *this;
return ret += rhs;
}
auto operator+(const formal_power_series &rhs) const {
auto ret = *this;
return ret += rhs;
}
auto &operator-=(const formal_power_series &rhs) {
if (data_.size() < rhs.size()) data_.resize(rhs.size());
for (int i = 0; i < rhs.size(); ++i) data_[i] -= rhs[i];
return *this;
}
auto &operator-=(T rhs) {
data_[0] -= rhs;
return *this;
}
auto operator-(T rhs) const {
auto ret = *this;
return ret -= rhs;
}
auto operator-(const formal_power_series &rhs) const {
auto ret = *this;
return ret -= rhs;
}
auto operator-() const {
auto ret = *this;
for (auto &x : ret) x = -x;
return ret;
}
auto &operator*=(const formal_power_series &rhs) {
data_ = convolve(data_, rhs.data_);
return *this;
}
auto operator*(const formal_power_series &rhs) const {
auto ret = convolve(data_, rhs.data_);
return formal_power_series(ret);
}
auto &operator*=(T rhs) {
for (auto &x : data_) x *= rhs;
return *this;
}
auto operator*(T rhs) const {
auto ret = *this;
return ret *= rhs;
}
auto differentiate() const {
const int n = data_.size();
std::vector<T> ret(n - 1);
for (int i = 0; i < n - 1; ++i) {
ret[i] = data_[i + 1] * (i + 1);
}
return formal_power_series(ret);
}
auto integrate() const {
const int n = data_.size();
std::vector<T> ret(n + 1), invs(n + 1, 1);
const int p = T::mod();
for (int i = 2; i <= n; ++i) invs[i] = -invs[p % i] * (p / i);
for (int i = 0; i < n; ++i) {
ret[i + 1] = data_[i] * invs[i + 1];
}
return formal_power_series(ret);
}
auto inv() const {
assert(data_[0] != 0);
const int n = data_.size();
int t = 1;
std::vector<T> ret = {data_[0].inv()};
ret.reserve(n * 2);
while (t <= n * 2) {
std::vector<T> c(data_.begin(), data_.begin() + std::min(t, n));
auto a = convolve(ret, ret, true);
if ((int) a.size() > t) a.resize(t);
c = convolve(c, a);
if ((int) c.size() > t) c.resize(t);
if ((int) ret.size() > t) ret.resize(t);
for (int i = 0; i < (int) ret.size(); ++i) ret[i] = ret[i] * 2;
if (ret.size() < c.size()) ret.resize(std::min<int>(c.size(), t));
for (int i = 0; i < (int) c.size(); ++i) {
ret[i] -= c[i];
}
t <<= 1;
}
ret.resize(n);
return formal_power_series(ret);
}
auto log() const {
assert(data_[0] == 1);
const int n = data_.size();
auto ret = (differentiate() * inv()).integrate();
ret.resize(n);
return ret;
}
auto exp() const {
const int n = data_.size();
int t = 1;
formal_power_series b({1});
while (t <= n * 2) {
t <<= 1;
auto temp = b.log();
temp.resize(t);
for (int i = 0; i < t; ++i) temp[i] = -temp[i];
temp[0] += 1;
for (int i = 0; i < std::min(t, n); ++i) temp[i] += data_[i];
b = b * temp;
b.resize(t);
}
b.resize(n);
return b;
}
auto shift(int64_t k) const {
const int64_t n = data_.size();
formal_power_series ret(n);
if (k >= 0) {
for (int64_t i = k; i < n; ++i) {
ret[i] = data_[i - k];
}
} else {
for (int64_t i = 0; i < n + k; ++i) {
ret[i] = data_[i - k];
}
}
return ret;
}
auto pow(int64_t M) const {
assert(M >= 0);
const int n = data_.size();
int k = 0;
for (; k < n; ++k) {
if (data_[k] != 0) {
break;
}
}
if (k >= n) return *this;
T a = data_[k];
formal_power_series ret = *this;
ret = (ret.shift(-k)) * a.inv();
ret = (ret.log() * (T) M).exp();
ret = (ret * a.pow(M)).shift(M * k);
return ret;
}
std::optional<formal_power_series> sqrt() const;
};
} // namespace haar_lib
#line 2 "Mylib/Number/Mod/mod_sqrt.cpp"
#include <optional>
#include <random>
#line 2 "Mylib/Number/Mod/mod_pow.cpp"
#include <cstdint>
namespace haar_lib {
constexpr int64_t mod_pow(int64_t n, int64_t p, int64_t m) {
int64_t ret = 1;
while (p > 0) {
if (p & 1) (ret *= n) %= m;
(n *= n) %= m;
p >>= 1;
}
return ret;
}
} // namespace haar_lib
#line 5 "Mylib/Number/Mod/mod_sqrt.cpp"
namespace haar_lib {
std::optional<int64_t> mod_sqrt(int64_t a, int64_t p) {
if (p == 2) return a % 2;
if (a == 0) return 0;
int64_t b = mod_pow(a, (p - 1) / 2, p);
if (b == p - 1) return {};
if (p % 4 == 3) return mod_pow(a, (p + 1) / 4, p);
int64_t q = p - 1, s = 0;
while (q % 2 == 0) {
q /= 2;
s += 1;
}
static std::mt19937_64 rand(time(0));
std::uniform_int_distribution<> dist(0, p - 1);
int64_t z;
while (1) {
z = dist(rand);
if (mod_pow(z, (p - 1) / 2, p) == p - 1) break;
}
int64_t m = s;
int64_t c = mod_pow(z, q, p);
int64_t t = mod_pow(a, q, p);
int64_t r = mod_pow(a, (q + 1) / 2, p);
while (1) {
if (t == 0) return 0;
if (t == 1) return r;
int i = 1;
for (int64_t T = t; i < m; ++i) {
(T *= T) %= p;
if (T == 1) break;
}
int64_t b = mod_pow(c, 1LL << (m - i - 1), p);
m = i;
c = b * b % p;
(t *= b * b % p) %= p;
(r *= b) %= p;
}
}
} // namespace haar_lib
#line 4 "Mylib/Math/fps_sqrt.cpp"
namespace haar_lib {
template <typename T, const auto &convolve>
auto formal_power_series<T, convolve>::sqrt() const -> std::optional<formal_power_series<T, convolve>> {
constexpr int mod = value_type::mod();
const int n = data_.size();
int k = 0;
for (; k < n; ++k)
if (data_[k] != 0) break;
if (k >= n) return *this;
if (k % 2 != 0) return {};
auto x = mod_sqrt((int64_t) data_[k], mod);
if (not x) return {};
const int m = n - k;
auto it = data_.begin() + k;
formal_power_series ret({*x});
int t = 1;
while (t <= m * 2) {
formal_power_series f(std::vector<T>(it, it + std::min(t, m)));
ret.resize(t);
f.resize(t);
ret = (ret + f * ret.inv()) * T(2).inv();
t <<= 1;
}
ret.resize(n);
ret = ret.shift(k / 2);
return ret;
}
} // namespace haar_lib
Back to top page
Mod pow
(Mylib/Number/Mod/mod_pow.cpp)