多项式开根 11E
2023 年 1 月 24 日
目录
多项式开根 11E
#ifndef ALGO_MATH_POLY_SQRT11ENT
#define ALGO_MATH_POLY_SQRT11ENT
#include "poly-def.hpp"
template <class ModT>
AVec<ModT> poly_sqrt_11E(std::span<const ModT> self, u32 m, const ModT &x0) {
u32 n = std::bit_ceil(m);
AVec<ModT> x(n), g(n), ng(n);
x[0] = x0;
if (n == 1)
return x;
ng[0] = g[0] = x[0].inv();
x[1] = (self[1] * g[0]).shift2();
ntt<ModT>({ng.begin(), 2});
for (u32 t = 2; t < n; t *= 2) {
AVec<ModT> f(t * 2), nf(t);
std::copy_n(x.begin(), t, nf.begin());
ntt<ModT>(nf); // 1E
std::copy_n(nf.begin(), t, f.begin());
dot<ModT>(nf, ng);
intt<ModT>(nf); // 1E
std::fill_n(nf.begin(), t / 2, 0);
ntt<ModT>(nf); // 1E
dot<ModT>(nf, ng);
intt<ModT>(nf); // 1E
vectorization_2(t / 2, g.data() + t / 2, nf.data() + t / 2, [](auto &gi, auto nfi) {
gi = -nfi;
});
dot<ModT>({f.begin(), t}, f);
intt<ModT>({f.begin(), t}); // 1E
for (u32 i = t; i < std::min<u32>(self.size(), t * 2); ++i)
f[i] = self[i - t] + self[i] - f[i - t];
std::fill_n(f.begin(), t, 0);
std::copy_n(g.begin(), t, ng.begin());
ntt<ModT>(f); // 2E
ntt<ModT>({ng.begin(), t * 2}); // 2E
dot<ModT>(f, ng);
intt<ModT>(f); // 2E
vectorization_2(t, x.data() + t, f.data() + t, [](auto &xi, auto fi) {
xi = fi.shift2();
});
}
return x.resize(m), x;
}
#endif // ALGO_MATH_POLY_SQRT11ENT