cp-book
ecnerwala's competitive programming library
🧪 fft/poly.test.cpp
Depends on
fft/common.hpp- fft/core.hpp
- fft/engine.hpp
- fft/engines/crt.hpp
- fft/engines/ntt.hpp
fft/engines/real.hpp- fft/engines/split.hpp
- fft/multiply.hpp
- fft/poly.hpp
- fft/series.hpp
- fft/series_core.hpp
- 🧪 fft/test_util.test.hpp
- modnum.hpp
Code
#include <bits/stdc++.h>
#include <catch2/catch_test_macros.hpp>
#include <catch2/catch_get_random_seed.hpp>
#include "fft/poly.hpp"
#include "fft/series.hpp"
#include "fft/test_util.test.hpp"
#include "modnum.hpp"
namespace ecnerwala {
namespace fft {
using namespace std;
TEST_CASE("poly::form evaluation and transposed multiplication", "[fft]") {
using num = modnum<998244353>;
using E = engines::ntt<num>;
mt19937 mt(Catch::getSeed());
int n = 40;
num z = num(mt());
auto f = poly::form<E>::polynomial_evaluation(z, n);
vector<num> sv(30);
fill_rnd(sv, mt);
poly::vec<E> s((span<const num>(sv)));
REQUIRE(f(s) == s(z));
// f(S * q) == composed_with(q)(S)
vector<num> qv(11), s2v(n - 11);
fill_rnd(qv, mt);
fill_rnd(s2v, mt);
poly::vec<E> q((span<const num>(qv))), s2((span<const num>(s2v)));
auto fq = f.composed_with(q);
REQUIRE(fq(s2) == f(s2 * q));
// evaluation functional composed with q evaluates S * q at z
REQUIRE(fq(s2) == (s2 * q)(z));
// composing with a power series (living in 1/x) multiplies the storages,
// prefix-truncated back to length n
series::trunc<E> t(size_t(n), num{});
fill_rnd(t, mt);
auto ft = f.composed_with(t);
REQUIRE(ft.len() == n);
for (int j = 0; j < n; j++) {
num want{};
for (int d = 0; d <= j; d++) want += t[d] * f.rev_series()[j - d];
REQUIRE(ft.rev_series()[j] == want);
}
// an exact series may be any length; the tail beyond it is zero
series::exact<E> e(t.begin(), t.begin() + 11);
auto fe = f.composed_with(e);
REQUIRE(fe.len() == n);
for (int j = 0; j < n; j++) {
num want{};
for (int d = 0; d <= j && d < 11; d++) want += t[d] * f.rev_series()[j - d];
REQUIRE(fe.rev_series()[j] == want);
}
}
TEST_CASE("poly::vec reversed storage and series interop", "[fft]") {
using num = modnum<998244353>;
using E = engines::ntt<num>;
mt19937 mt(Catch::getSeed());
vector<num> pa(37), pb(23);
fill_rnd(pa, mt);
fill_rnd(pb, mt);
poly::vec<E> a((span<const num>(pa))), b((span<const num>(pb)));
// indexing is coefficient order, storage is reversed
REQUIRE(a[0] == pa[0]);
REQUIRE(a.leading() == pa[36]);
REQUIRE(a.rev_series()[0] == pa[36]);
REQUIRE(a.rev_series()[36] == pa[0]);
// products convolve the storage directly
poly::vec<E> p = a * b;
check_eq(vector<num>(p.begin(), p.end()), multiply_slow(pa, pb));
REQUIRE(square(a) == a * a);
num x = num(mt());
REQUIRE(p(x) == a(x) * b(x));
// +/- align at x^0 (the shared storage tail)
poly::vec<E> s = a + b, d = b - a;
REQUIRE(s.len() == 37);
for (int i = 0; i < 37; i++) REQUIRE(s[i] == pa[i] + (i < 23 ? pb[i] : num(0)));
for (int i = 0; i < 37; i++) REQUIRE(d[i] == (i < 23 ? pb[i] : num(0)) - pa[i]);
// multiplying by x^k appends zeros to the storage; coefficients shift up
poly::vec<E> g = a;
g.shift(2);
REQUIRE(g.len() == 39);
REQUIRE(g[0] == num(0));
REQUIRE(g[1] == num(0));
for (int i = 0; i < 37; i++) REQUIRE(g[i + 2] == pa[i]);
REQUIRE(g.rev_series().data()[0] == pa[36]);
// named conversions use the reversed convention and round-trip freely
const series::exact<E>& ra = a.rev_series();
REQUIRE(ra.len() == 37);
for (int i = 0; i < 37; i++) REQUIRE(ra[i] == pa[36 - i]);
REQUIRE(poly::vec<E>::from_rev_series(ra) == a);
// a poly::vec's natural-order coefficients as an exact series (for series products)
series::exact<E> xa(a.begin(), a.end());
REQUIRE(equal(xa.begin(), xa.end(), pa.begin(), pa.end()));
REQUIRE(a.unrev_series(10) == series::trunc<E>(pa.begin(), pa.begin() + 10));
// the storage transform serves transposed products: middle product against rev_series()
std::vector<num> vals(60);
fill_rnd(vals, mt);
series::cached_exact<E> cv(series::exact<E>(vals.begin(), vals.end()));
series::cached_exact<E> ca(a.rev_series());
auto mp = middle_product(cv, ca);
auto naive = [&](int j) {
num r{};
for (int t = 0; t < 37; t++) r += pa[t] * vals[j + t];
return r;
};
for (int j = 0; j < sz(mp); j++) REQUIRE(mp[size_t(j)] == naive(j));
}
TEST_CASE("poly::cached products", "[fft]") {
using num = modnum<998244353>;
using E = engines::ntt<num>;
mt19937 mt(Catch::getSeed());
vector<num> pa(37), pb(23);
fill_rnd(pa, mt);
fill_rnd(pb, mt);
poly::vec<E> a((span<const num>(pa))), b((span<const num>(pb)));
// poly::vec products return poly::cached, adopting the product transform
auto p = a * b;
static_assert(std::is_same_v<decltype(p), poly::cached<E>>);
REQUIRE(p.rev_series().cache().size() > 0);
poly::vec<E> pp = a * b; // naming the plain type moves out and drops the transform
check_eq(vector<num>(pp.begin(), pp.end()), multiply_slow(pa, pb));
// cached operands reuse and chain; results compare across representations
poly::cached<E> ca(a), cb(b);
REQUIRE(ca == a);
REQUIRE(ca * cb == p);
REQUIRE(ca * b == p);
auto sq = square(ca);
static_assert(std::is_same_v<decltype(sq), poly::cached<E>>);
REQUIRE(sq == a * a);
num x = num(mt());
REQUIRE(p(x) == ca(x) * cb(x));
// moving out drops down to a plain mutable poly::vec
poly::vec<E> q = std::move(ca);
REQUIRE(q == a);
}
TEST_CASE("poly::multipoint and poly::interpolate", "[fft]") {
using num = modnum<998244353>;
mt19937 mt(Catch::getSeed());
for (int n : {1, 2, 3, 8, 17, 40}) {
INFO("n = " << n);
vector<num> coeffs(n);
for (num& x : coeffs) { x = num(mt()); }
poly::vec<engines::ntt<num>> p((span<const num>(coeffs)));
vector<num> pts(n);
for (int i = 0; i < n; i++) pts[i] = num(1000 + i);
auto vals = poly::multipoint<engines::ntt<num>>(p, pts);
for (int i = 0; i < n; i++) {
REQUIRE(vals[i] == p(pts[i]));
}
auto rec = poly::interpolate<engines::ntt<num>>(pts, vals);
REQUIRE(rec == p);
}
}
}} // namespace ecnerwala::fft
#include <bits/stdc++.h>
#include <catch2/catch_test_macros.hpp>
#include <catch2/catch_get_random_seed.hpp>
#line 4 "src/fft/poly.test.cpp"
#line 2 "src/fft/poly.hpp"
#line 12 "src/fft/poly.hpp"
#line 2 "src/fft/series_core.hpp"
#line 12 "src/fft/series_core.hpp"
#line 2 "src/fft/multiply.hpp"
#line 10 "src/fft/multiply.hpp"
#line 2 "src/fft/engine.hpp"
#line 7 "src/fft/engine.hpp"
#line 2 "src/fft/common.hpp"
#line 8 "src/fft/common.hpp"
/**
* Author: Andrew He
* Source: http://neerc.ifmo.ru/trains/toulouse/2017/fft2.pdf
* Papers about accuracy: http://www.daemonology.net/papers/fft.pdf, http://www.cs.berkeley.edu/~fateman/papers/fftvsothers.pdf
* For integers rounding works if $(|a| + |b|)\max(a, b) < \mathtt{\sim} 10^9$, or in theory maybe $10^6$.
*
* Abstraction layers:
* fft_core<num> FFT itself and other ops on rings with 2^k-th roots of unity. We use bit-reversed indexing in the frequency domain.
*
* engines Engines for packing/unpacking arbitrary rings for convolution (the `engine` concept).
* Still expose (opaque) transform-domain objects for caching/fusion.
*
* multiply layer Wrappers for convolving bounded sequences: track length/truncation.
*
* value types series::vec<E, exact> - R[[x]]
* series::exact<E> - exact (finite-support) power series
* series::trunc<E> - truncated prefix of an (infinite) power series
*
* polynomials - R[x]. Under x -> 1/x a polynomial becomes a Laurent polynomial in 1/x;
* shifting by x^{deg P} (reversal) lands it in R[[x]], and we store that exact series.
* poly::vec<E> - polynomial type, supporting natural indexing
* poly::form<E> - finite-support linear forms, via the pairing <P, S> = [x^0] P(1/x) S(x)
* a linear form is one side of this pairing, applied to the other
*
* online_multiplier<E> - online (relaxed) multiplication of 2 sequences in n log^2 n time
* ap_sampled_poly<E> - a polynomial stored as its evaluations on an arithmetic progression
*/
namespace ecnerwala {
template<class T> int sz(T&& arg) { using std::size; return int(size(std::forward<T>(arg))); }
inline int nextPow2(int s) { return 1 << (s > 1 ? 32 - __builtin_clz(s-1) : 0); }
namespace fft {
using std::swap;
using std::vector;
using std::min;
using std::max;
// Reusable scratch buffers. Not thread-safe by default: this is deliberately plain
// static storage so single-threaded programs pay no TLS indirection; define
// ECNERWALA_FFT_POOL_STORAGE to `thread_local` for multithreaded use.
#ifndef ECNERWALA_FFT_POOL_STORAGE
#define ECNERWALA_FFT_POOL_STORAGE
#endif
template <typename T> struct buffer_pool {
static inline ECNERWALA_FFT_POOL_STORAGE std::vector<std::vector<T>> free_list;
struct handle {
std::vector<T> v;
explicit handle(int n) {
if (!free_list.empty()) {
v = std::move(free_list.back());
free_list.pop_back();
}
v.assign(n, T());
}
handle(const handle&) = delete;
handle& operator=(const handle&) = delete;
handle(handle&& o) noexcept : v(std::move(o.v)) {}
~handle() { if (v.capacity()) free_list.push_back(std::move(v)); }
T& operator[](int i) { return v[i]; }
operator std::span<T>() { return std::span<T>(v); }
std::span<T> span() { return std::span<T>(v); }
};
static handle get(int n) { return handle(n); }
};
/* namespace fft */ }
/* namespace ecnerwala */ }
#line 9 "src/fft/engine.hpp"
namespace ecnerwala::fft {
// ==== engine concept ====
// Output operations for the finish step to express arbitrary fusion into the output buffer.
struct assign_op { template <typename T> void operator()(T& d, T v) const { d = v; } };
struct add_op { template <typename T> void operator()(T& d, T v) const { d += v; } };
struct sub_op { template <typename T> void operator()(T& d, T v) const { d -= v; } };
struct add_twice_op { template <typename T> void operator()(T& d, T v) const { d += v + v; } };
// `engine` contract
// engine represents a way of packing/unpacking sequences over an arbitrary ring into FFT-style transforms.
// We expect transforms/products of transforms to be linear but potentially lossy/imprecise, so we'll track precision
// at compile-time as a template parameter.
//
// E::value_type The ring we operate over
// E::unit_scale 0 or 1 depending on whether there's error that can accumulate
// E::commutative A marker for whether the ring is commutative
//
// transformed_t<A> The transform of a sequence. This object owns its data buffer.
// product_t<A> The product of 2 transforms. May equal transformed_t, particularly when unit_scale = 0.
//
// transformed alias for transformed_t<unit_scale>
// product alias for product_t<unit_scale>
//
// The basic multiplication API is
// transform(span<const value_type> in, int n) -> transformed_t<unit_scale>
// mul(transformed_t<A>, transformed_t<B>, int n) -> product_t<A*B>
// mul2(a1, b1, a2, b2, int n) -> product_t<A1*B1 + A2*B2>, computing a1*b1 + a2*b2 in one pass
// finish(product_t<A>, span<value_type>& out, Op) -> void
//
// Input span can be length up to 2n.
// Output spans can be length up to n; only the prefix that exists is filled.
// finish applies Op exactly once per out element, in index order, to
// value_type targets (so ops may be stateful).
// Transforms can be longer than necessary, and only the relevant prefix is used.
//
// For non-exact engines, there's some subtlety in whether we wrap before or after packing.
// We will choose to wrap *after* packing, which hurts error bounds but makes the prefix condition more uniform.
//
// Additionally, we have APIs to take advantage of linearity in both transformed and product space:
// add(transformed_t<A>, transformed_t<B>) -> transformed_t<A+B>
// add(product_t<K1>, product_t<K2>) -> product_t<K1+K2>
//
// Finally, we expose some additional fast-transform optimization paths.
// extend_to only operates on transformed_t<unit_scale>; the others are scale-generic.
// downsample is also defined on product_t (halving before finish saves inverse-transform work).
// extend_to build (if empty) or grow a transform to size m by repeated doubling; feed the same coefficients (sz <= 2m) every time
// (each doubling step reads only the coefficients that fit, so zero-padded buffers are fine:
// coefficients past twice the existing transform's size must be zero, or it couldn't be a prefix)
// downsample compute the half-sized transform/product of just the even (odd = false) or odd terms of the input
// negate_arg size n transform of A(-x)
template <typename E>
concept engine = requires(
std::span<const typename E::value_type> in,
std::span<typename E::value_type> out,
typename E::transformed& t,
const typename E::transformed& ct,
typename E::product& p,
const typename E::product& cp,
int n
) {
typename E::value_type;
{ E::transform(in, n) } -> std::same_as<typename E::transformed>;
{ ct.size() } -> std::same_as<int>;
E::extend_to(t, n, in);
{ E::downsample(ct, n, false) } -> std::same_as<typename E::transformed>;
{ E::downsample(cp, n, false) } -> std::same_as<typename E::product>;
{ E::negate_arg(ct, n) } -> std::same_as<typename E::transformed>;
{ E::mul(ct, ct, n) } -> std::same_as<typename E::product>;
{ E::sq(ct, n) } -> std::same_as<typename E::product>;
{ E::mul2(ct, ct, ct, ct, n) } -> std::same_as<typename E::template product_t<2 * E::unit_scale>>;
E::finish(std::move(p), out);
E::finish(std::move(p), out, add_op{});
E::finish(E::add(std::move(p), std::move(p)), out);
{ E::add(E::transform(in, n), ct) } -> std::same_as<typename E::template transformed_t<2 * E::unit_scale>>;
{ E::add(std::move(p), std::move(p)) } -> std::same_as<typename E::template product_t<2 * E::unit_scale>>;
requires std::same_as<std::remove_cvref_t<decltype(E::commutative)>, bool>;
requires std::same_as<std::remove_cvref_t<decltype(E::unit_scale)>, int>;
};
// Constrains two engine-parameterized value types to share the same engine.
template <typename A, typename B>
concept same_engine = std::same_as<typename A::engine_t, typename B::engine_t>;
// short spelling for E::transformed at use sites
template <engine E> using transformed = typename E::transformed;
/* namespace ecnerwala::fft */ }
#line 12 "src/fft/multiply.hpp"
namespace ecnerwala::fft {
// ==== multiply layer ====
// These are free functions to convolve spans.
//
// The interfaces will typically take input spans, an output span, and an Op representing how to fold the result into the output.
// Output spans may alias one of the input spans.
// Output spans may be shorter than expected; the output will just be truncated.
//
// Some functions may also take E::transformed& objects associated with the input
// spans. These will be lazily filled (see E::extend_to) and used if available.
// Circular convolution mod n (power of 2)
template <engine E, typename Op = assign_op>
void multiply_circular(std::span<const typename E::value_type> a, std::span<const typename E::value_type> b,
std::span<typename E::value_type> out, int n, Op op = {}) {
assert(!(n & (n-1)));
auto ta = E::transform(a, n);
auto tb = E::transform(b, n);
E::finish(E::mul(ta, tb, n), out, op);
}
template <engine E, typename Op = assign_op>
void square_circular(std::span<const typename E::value_type> a, std::span<typename E::value_type> out, int n, Op op = {}) {
assert(!(n & (n-1)));
auto ta = E::transform(a, n);
E::finish(E::sq(ta, n), out, op);
}
namespace detail {
// Arrays of length 2^k + 1 are somewhat common, so we will optimize them by
// multiplying mod 2^k, and fixing up the leading coefficient.
// Helpers to detect and perform this optimization.
struct conv_size { int n; bool cut; };
inline conv_size conv_size_for(int s) {
int n = nextPow2(s);
bool cut = (n == 2 * (s - 1));
return {cut ? n / 2 : n, cut};
}
// Call op while lazily applying the correction if necessary
template <typename T, typename Op>
void emit_linear(std::span<T> buf, int n, int s, bool cut, T c0, std::span<T> out, Op op) {
T cn{};
if (cut) {
cn = buf[0] - c0;
buf[0] = c0;
}
int lim = min(sz(out), min(s, n));
for (int i = 0; i < lim; i++) op(out[i], buf[i]);
if (cut && sz(out) >= s) op(out[s-1], cn);
}
// Applies op, diverting the wrapped leading coefficient of a cut product:
// out[0] receives c0 and the wraparound term is captured into cn for the
// caller to emit at out[s-1].
template <typename T, typename Op>
struct cut_op {
Op op;
T* out0;
T c0;
T& cn;
void operator()(T& x, T v) const {
if (&x == out0) { cn = v - c0; v = c0; }
op(x, v);
}
};
// finish + emit_linear fused: write the finished product directly into out,
// applying the cut correction in place.
template <engine E, typename P, typename Op = assign_op>
void finish_linear(
P&& p, int n, int s, bool cut,
typename E::value_type c0, std::span<typename E::value_type> out, Op op = {}
) {
using T = typename E::value_type;
if (sz(out) == 0) return;
int lim = min(sz(out), min(s, n));
if (!cut) {
E::finish(std::move(p), out.subspan(0, lim), op);
} else {
T cn{};
E::finish(std::move(p), out.subspan(0, lim), cut_op<T, Op>{op, &out[0], c0, cn});
if (sz(out) >= s) op(out[s-1], cn);
}
}
}
template <engine E, typename Op = assign_op>
void multiply(std::span<const typename E::value_type> a, std::span<const typename E::value_type> b,
std::span<typename E::value_type> out, Op op = {}) {
using T = typename E::value_type;
if (sz(a) == 0 || sz(b) == 0) return;
int s = sz(a) + sz(b) - 1;
auto [n, cut] = detail::conv_size_for(s);
T c0 = a[0] * b[0];
auto buf = buffer_pool<T>::get(n);
multiply_circular<E>(a, b, buf.span(), n);
detail::emit_linear<T>(buf.span(), n, s, cut, c0, out, op);
}
template <engine E, typename Op = assign_op>
void multiply(std::span<const typename E::value_type> a, transformed<E>& ta,
std::span<const typename E::value_type> b, transformed<E>& tb,
std::span<typename E::value_type> out, Op op = {}) {
using T = typename E::value_type;
if (sz(a) == 0 || sz(b) == 0) return;
int s = sz(a) + sz(b) - 1;
auto [n, cut] = detail::conv_size_for(s);
T c0 = a[0] * b[0];
E::extend_to(ta, n, a);
E::extend_to(tb, n, b);
detail::finish_linear<E>(E::mul(ta, tb, n), n, s, cut, c0, out, op);
}
template <engine E, typename Op = assign_op>
void multiply_add2(std::span<const typename E::value_type> a1, transformed<E>& ta1,
std::span<const typename E::value_type> b1, transformed<E>& tb1,
std::span<const typename E::value_type> a2, transformed<E>& ta2,
std::span<const typename E::value_type> b2, transformed<E>& tb2,
std::span<typename E::value_type> out, Op op = {}) {
using T = typename E::value_type;
assert(sz(a1) > 0 && sz(b1) > 0 && sz(a2) > 0 && sz(b2) > 0);
int s = sz(a1) + sz(b1) - 1;
assert(sz(a2) + sz(b2) - 1 == s);
auto [n, cut] = detail::conv_size_for(s);
T c0 = a1[0] * b1[0] + a2[0] * b2[0];
E::extend_to(ta1, n, a1); E::extend_to(tb1, n, b1);
E::extend_to(ta2, n, a2); E::extend_to(tb2, n, b2);
detail::finish_linear<E>(E::mul2(ta1, tb1, ta2, tb2, n), n, s, cut, c0, out, op);
}
// As multiply_add2, but also outputs the summed pointwise product as a reusable
// transform of the (full-length) result, like multiply_cached.
template <engine E>
void multiply_add2_cached(
std::span<const typename E::value_type> a1, transformed<E>& ta1,
std::span<const typename E::value_type> b1, transformed<E>& tb1,
std::span<const typename E::value_type> a2, transformed<E>& ta2,
std::span<const typename E::value_type> b2, transformed<E>& tb2,
std::vector<typename E::value_type>& coeffs, transformed<E>& t) {
using T = typename E::value_type;
assert(sz(a1) > 0 && sz(b1) > 0 && sz(a2) > 0 && sz(b2) > 0);
int s = sz(a1) + sz(b1) - 1;
assert(sz(a2) + sz(b2) - 1 == s);
coeffs.assign(size_t(s), T{});
t = transformed<E>{};
if constexpr (std::same_as<typename E::product, transformed<E>>) {
auto [n, cut] = detail::conv_size_for(s);
T c0 = a1[0] * b1[0] + a2[0] * b2[0];
E::extend_to(ta1, n, a1); E::extend_to(tb1, n, b1);
E::extend_to(ta2, n, a2); E::extend_to(tb2, n, b2);
auto p = E::mul2(ta1, tb1, ta2, tb2, n);
auto tp = p;
detail::finish_linear<E>(std::move(p), n, s, cut, c0, std::span<T>(coeffs));
t = std::move(tp);
} else {
multiply_add2<E>(a1, ta1, b1, tb1, a2, ta2, b2, tb2, std::span<T>(coeffs));
}
}
// This helper also accepts an output transform which will be populated if it is cheap to do so
template <engine E>
void multiply_cached(std::span<const typename E::value_type> a, transformed<E>& ta,
std::span<const typename E::value_type> b, transformed<E>& tb,
std::vector<typename E::value_type>& coeffs, transformed<E>& t) {
using T = typename E::value_type;
coeffs.assign(size_t(sz(a) && sz(b) ? sz(a) + sz(b) - 1 : 0), T{});
t = transformed<E>{};
if (coeffs.empty()) return;
int s = sz(coeffs);
if constexpr (std::same_as<typename E::product, transformed<E>>) {
auto [n, cut] = detail::conv_size_for(s);
T c0 = a[0] * b[0];
E::extend_to(ta, n, a);
E::extend_to(tb, n, b);
auto p = E::mul(ta, tb, n);
auto tp = p;
detail::finish_linear<E>(std::move(p), n, s, cut, c0, std::span<T>(coeffs));
t = std::move(tp);
} else {
multiply<E>(a, ta, b, tb, std::span<T>(coeffs));
}
}
template <engine E, typename Op = assign_op>
void square(std::span<const typename E::value_type> a, std::span<typename E::value_type> out, Op op = {}) {
using T = typename E::value_type;
if (sz(a) == 0) return;
int s = 2 * sz(a) - 1;
auto [n, cut] = detail::conv_size_for(s);
T c0 = a[0] * a[0];
auto buf = buffer_pool<T>::get(n);
square_circular<E>(a, buf.span(), n);
detail::emit_linear<T>(buf.span(), n, s, cut, c0, out, op);
}
template <engine E, typename Op = assign_op>
void square(std::span<const typename E::value_type> a, transformed<E>& ta,
std::span<typename E::value_type> out, Op op = {}) {
using T = typename E::value_type;
if (sz(a) == 0) return;
int s = 2 * sz(a) - 1;
auto [n, cut] = detail::conv_size_for(s);
T c0 = a[0] * a[0];
E::extend_to(ta, n, a);
detail::finish_linear<E>(E::sq(ta, n), n, s, cut, c0, out, op);
}
// As square, but also outputs the pointwise product as a reusable transform of
// the result (empty when the engine's product isn't a transform).
template <engine E>
void square_cached(std::span<const typename E::value_type> a, transformed<E>& ta,
std::vector<typename E::value_type>& coeffs, transformed<E>& t) {
using T = typename E::value_type;
coeffs.assign(size_t(sz(a) ? 2 * sz(a) - 1 : 0), T{});
t = transformed<E>{};
if (coeffs.empty()) return;
int s = sz(coeffs);
if constexpr (std::same_as<typename E::product, transformed<E>>) {
auto [n, cut] = detail::conv_size_for(s);
T c0 = a[0] * a[0];
E::extend_to(ta, n, a);
auto p = E::sq(ta, n);
auto tp = p;
detail::finish_linear<E>(std::move(p), n, s, cut, c0, std::span<T>(coeffs));
t = std::move(tp);
} else {
square<E>(a, ta, std::span<T>(coeffs));
}
}
template <engine E> vector<typename E::value_type> multiply(
const vector<typename E::value_type>& a, const vector<typename E::value_type>& b) {
using T = typename E::value_type;
if (sz(a) == 0 || sz(b) == 0) return {};
vector<T> r(sz(a) + sz(b) - 1);
multiply<E>(std::span<const T>(a), std::span<const T>(b), std::span<T>(r));
return r;
}
template <engine E> vector<typename E::value_type> square(const vector<typename E::value_type>& a) {
using T = typename E::value_type;
if (sz(a) == 0) return {};
vector<T> r(2 * sz(a) - 1);
square<E>(std::span<const T>(a), std::span<T>(r));
return r;
}
namespace detail {
// emit_linear but for middle_product
template <typename T, typename Op>
void emit_middle(std::span<T> buf, bool cut, int la, int lb, T c0, T ctop, std::span<T> out, Op op) {
int m = la - lb + 1;
T cn{};
if (cut) {
cn = buf[0] - c0; // for lb == 1 these coincide: slot 0 = c_0 + c_n and ctop = c_n
buf[lb - 1] -= ctop;
}
int lim = min(sz(out), cut ? m - 1 : m);
for (int t = 0; t < lim; t++) op(out[t], buf[lb - 1 + t]);
if (cut && sz(out) >= m) op(out[m-1], cn);
}
}
// Middle product (the transposed multiplication): takes only coefficients of a * b which include terms from all of b.
// Must have len(a) >= len(b)
template <engine E, typename Op = assign_op>
void middle_product(std::span<const typename E::value_type> a, std::span<const typename E::value_type> b,
std::span<typename E::value_type> out, Op op = {}) {
using T = typename E::value_type;
if (sz(a) == 0 || sz(b) == 0) return;
assert(sz(a) >= sz(b));
if (sz(a) == sz(b)) {
T r{};
for (int i = 0; i < sz(a); i++) {
r += a[i] * b[sz(b) - 1 - i];
}
if (sz(out) > 0) op(out[0], r);
return;
}
auto [n, cut] = detail::conv_size_for(sz(a));
auto buf = buffer_pool<T>::get(n);
multiply_circular<E>(a, b, buf.span(), n);
detail::emit_middle<T>(buf.span(), cut, sz(a), sz(b),
a[0] * b[0], a[sz(a) - 1] * b[sz(b) - 1], out, op);
}
// TODO: Let's decide whether to keep vector<> returning forms or not; this
// largely depends on whether we think these functions are a public interface or
// merely convenience for value type implementors.
template <engine E> vector<typename E::value_type> middle_product(
std::span<const typename E::value_type> a, std::span<const typename E::value_type> b) {
using T = typename E::value_type;
if (sz(a) == 0 || sz(b) == 0) return {};
assert(sz(a) >= sz(b));
vector<T> r(size_t(sz(a) - sz(b) + 1));
middle_product<E>(a, b, std::span<T>(r));
return r;
}
template <engine E, typename Op = assign_op>
void middle_product(std::span<const typename E::value_type> a, transformed<E>& ta,
std::span<const typename E::value_type> b, transformed<E>& tb,
std::span<typename E::value_type> out, Op op = {}) {
using T = typename E::value_type;
if (sz(a) == 0 || sz(b) == 0) return;
assert(sz(a) >= sz(b));
if (sz(a) == sz(b)) {
T r{};
for (int i = 0; i < sz(a); i++) {
r += a[i] * b[sz(b) - 1 - i];
}
if (sz(out) > 0) op(out[0], r);
return;
}
auto [n, cut] = detail::conv_size_for(sz(a));
E::extend_to(ta, n, a);
E::extend_to(tb, n, b);
auto buf = buffer_pool<T>::get(n);
E::finish(E::mul(ta, tb, n), buf.span());
detail::emit_middle<T>(buf.span(), cut, sz(a), sz(b),
a[0] * b[0], a[sz(a) - 1] * b[sz(b) - 1], out, op);
}
template <engine E>
vector<typename E::value_type> middle_product(std::span<const typename E::value_type> a, transformed<E>& ta,
std::span<const typename E::value_type> b, transformed<E>& tb) {
using T = typename E::value_type;
if (sz(a) == 0 || sz(b) == 0) return {};
assert(sz(a) >= sz(b));
vector<T> r(size_t(sz(a) - sz(b) + 1));
middle_product<E>(a, ta, b, tb, std::span<T>(r));
return r;
}
/* namespace ecnerwala::fft */ }
#line 14 "src/fft/series_core.hpp"
// ==== value types ====
namespace ecnerwala::series {
// A series is either exact (a finite series, R[x] sitting inside R[[x]]: the
// length is just the support bound) or trunc (a known prefix of an infinite
// series: the length is the precision, and products truncate to it).
// Non-owning view of power series coefficients: the span pattern (contiguous
// window + series semantics), borrowed from an owning series-like type.
template <fft::engine E, bool exact_>
struct span {
using T = typename E::value_type;
using engine_t = E;
static constexpr bool exact_v = exact_;
span() = default;
explicit span(std::span<const T> s_) : s(s_) {}
// exact -> trunc is implicit, trunc -> exact is explicit
template <bool oe> requires (oe != exact_)
explicit(oe < exact_) span(span<E, oe> o) : s(o.coeffs()) {}
int len() const { return sz(s); }
const T& operator[](int i) const { return s[size_t(i)]; }
auto begin() const { return s.begin(); }
auto end() const { return s.end(); }
// engine primitives borrow through std::span's range constructor
std::span<const T> coeffs() const { return s; }
// the first n coefficients; requires n <= len().
// Widening past len() is explicit: see with_len.
span first(int n) const {
assert(n <= len());
return span(s.first(size_t(n)));
}
private:
std::span<const T> s;
};
// `vec` represents both exact (finite) power series (R[x]) and prefixes of infinite power series (R[[x]]), depending on the flag.
// `exact` and `trunc` are aliases.
//
// Operators here are typically permissive: they will accept combinations of unequal types and lengths.
template <fft::engine E, bool exact_>
struct vec : public std::vector<typename E::value_type> {
using T = typename E::value_type;
using engine_t = E;
static constexpr bool exact_v = exact_;
using std::vector<T>::vector;
// a free const borrow of the coefficients: implicit
operator span<E, exact_>() const {
return span<E, exact_>(std::span<const T>(*this));
}
// exact -> trunc is implicit, trunc -> exact is explicit
template <bool oe> requires (oe != exact_)
explicit(oe < exact_) vec(const vec<E, oe>& p) : std::vector<T>(p) {}
template <bool oe> requires (oe != exact_)
explicit(oe < exact_) vec(vec<E, oe>&& p) : std::vector<T>(std::move(p)) {}
// adopt a plain coefficient vector
explicit vec(std::vector<T> v) : std::vector<T>(std::move(v)) {}
// materialize an owned copy of any borrowed series, of either exactness
explicit vec(span<E, exact_> s) : std::vector<T>(s.begin(), s.end()) {}
explicit vec(span<E, !exact_> s) : std::vector<T>(s.begin(), s.end()) {}
span<E, exact_> first(int n) const { return span<E, exact_>(*this).first(n); }
int len() const {
return int(this->size());
}
int degree() const requires (exact_) {
return len() - 1;
}
void extend(int sz) {
assert(sz >= len());
this->resize(sz);
}
void shrink(int sz) {
assert(sz <= len());
this->resize(sz);
}
// multiply by x^n within the fixed precision window
void shift_trunc(int n = 1) requires (!exact_) {
assert(n >= 0 && n <= len());
std::rotate(this->begin(), this->end()-n, this->end());
std::fill(this->begin(), this->begin()+n, T(0));
}
// divide by x^n and 0-pad within the fixed precision window
void unshift_trunc(int n = 1) requires (!exact_) {
assert(n >= 0 && n <= len());
std::fill(this->begin(), this->begin()+n, T(0));
std::rotate(this->begin(), this->begin()+n, this->end());
}
// in-place forms require that the result's exactness/length must equal this operand's
template <bool oe> requires (exact_ <= oe)
vec& operator += (const vec<E, oe>& o) {
if constexpr (exact_) { if (o.len() > len()) this->resize(o.len()); }
else if constexpr (!oe) { if (o.len() < len()) this->resize(o.len()); }
for (int i = 0; i < std::min(len(), o.len()); i++) {
(*this)[i] += o[i];
}
return *this;
}
template <bool oe> requires (exact_ <= oe)
vec& operator -= (const vec<E, oe>& o) {
if constexpr (exact_) { if (o.len() > len()) this->resize(o.len()); }
else if constexpr (!oe) { if (o.len() < len()) this->resize(o.len()); }
for (int i = 0; i < std::min(len(), o.len()); i++) {
(*this)[i] -= o[i];
}
return *this;
}
vec& operator *= (const T& n) {
for (auto& v : *this) v *= n;
return *this;
}
friend vec operator * (const vec& a, const T& n) {
vec r(a.size());
for (int i = 0; i < a.len(); i++) {
r[i] = a[i] * n;
}
return r;
}
friend vec operator * (const T& n, const vec& a) {
vec r(a.size());
for (int i = 0; i < a.len(); i++) {
r[i] = n * a[i];
}
return r;
}
vec& operator *= (const vec& o) {
return *this = (*this) * o;
}
};
template <fft::engine E> using exact = vec<E, true>;
template <fft::engine E> using trunc = vec<E, false>;
// Series-like concepts: the binary operators below are written once as constrained
// templates and dispatch on which memoized transforms an operand carries.
// A series-like type exposes its engine/exactness and its coefficients as a
// span borrow of exactly len() coefficients; cached wrappers additionally
// expose their transform caches (filling them is logically const).
template <typename S>
concept like = fft::engine<typename S::engine_t> && requires(const S& s, int i) {
{ S::exact_v } -> std::convertible_to<bool>;
{ s.len() } -> std::same_as<int>;
{ s[i] } -> std::convertible_to<const typename S::engine_t::value_type&>;
// borrows straight into the engine primitives
{ std::span<const typename S::engine_t::value_type>(s) };
// and into the series layer's own span, keeping the exactness tag
requires std::convertible_to<const S&, span<typename S::engine_t, S::exact_v>>;
// first(n): the first n coefficients, borrowed; requires n <= len().
// The result is itself like (concepts can't self-reference).
// Cached types keep a cache in the result only when it still serves the whole borrow.
{ s.first(i) } -> std::convertible_to<span<typename S::engine_t, S::exact_v>>;
};
template <typename S>
concept exact_like = like<S> && S::exact_v;
template <typename S>
concept trunc_like = like<S> && !S::exact_v;
// carries one extendable transform of the whole coefficient sequence
template <typename S>
concept has_cache = like<S> && requires(const S& s) {
{ s.cache() } -> std::same_as<fft::transformed<typename S::engine_t>&>;
};
template <fft::engine E, bool exact_>
struct maybe_cached;
// A borrowed series paired with the transform serving it: the
// normalized operand form fed to the cached fft:: entry points. Models has_cache.
template <fft::engine E, bool exact_>
struct cached_span {
using engine_t = E;
static constexpr bool exact_v = exact_;
span<E, exact_> s;
std::reference_wrapper<fft::transformed<E>> f;
cached_span(span<E, exact_> s_, fft::transformed<E>& f_) : s(s_), f(f_) {}
// exact -> trunc is implicit, trunc -> exact is explicit
template <bool oe> requires (oe != exact_)
explicit(oe < exact_) cached_span(cached_span<E, oe> o) : s(span<E, exact_>(o.s)), f(o.f) {}
int len() const { return s.len(); }
const typename E::value_type& operator[](int i) const { return s[i]; }
operator std::span<const typename E::value_type>() const { return s.coeffs(); }
operator span<E, exact_>() const { return s; }
maybe_cached<E, exact_> first(int n) const;
fft::transformed<E>& cache() const { return f; }
};
// carries a whole-sequence cache only sometimes, decided at runtime
template <typename S>
concept has_cache_opt = like<S> && requires(const S& s) {
{ s.cache_opt() } -> std::same_as<std::optional<std::reference_wrapper<fft::transformed<typename S::engine_t>>>>;
};
namespace detail {
// the operand's whole cache, if it carries one
template <like S>
std::optional<std::reference_wrapper<fft::transformed<typename S::engine_t>>> cache_of(const S& s) {
if constexpr (has_cache<S>) return s.cache();
else if constexpr (has_cache_opt<S>) return s.cache_opt();
else return std::nullopt;
}
/* namespace detail */ }
// A borrowed series which may carry the transform serving it: the runtime
// counterpart of cached_span in the borrow hierarchy
// prefix_cached/cached -> maybe_cached/cached_span -> span.
template <fft::engine E, bool exact_>
struct maybe_cached {
using T = typename E::value_type;
using engine_t = E;
static constexpr bool exact_v = exact_;
span<E, exact_> s;
std::optional<std::reference_wrapper<fft::transformed<E>>> f;
explicit maybe_cached(span<E, exact_> s_) : s(s_) {}
maybe_cached(span<E, exact_> s_, fft::transformed<E>& f_) : s(s_), f(f_) {}
maybe_cached(cached_span<E, exact_> c) : s(c.s), f(c.f) {}
// borrow any like operand whole, taking along whatever cache it carries
template <like S> requires std::same_as<typename S::engine_t, E> && (S::exact_v == exact_)
maybe_cached(const S& o) : s(o), f(detail::cache_of(o)) {}
int len() const { return s.len(); }
const T& operator[](int i) const { return s[i]; }
operator std::span<const T>() const { return s.coeffs(); }
operator span<E, exact_>() const { return s; }
maybe_cached first(int n) const {
return n == len() ? *this : maybe_cached(s.first(n));
}
std::optional<std::reference_wrapper<fft::transformed<E>>> cache_opt() const { return f; }
};
template <fft::engine E, bool exact_>
maybe_cached<E, exact_> cached_span<E, exact_>::first(int n) const {
return maybe_cached<E, exact_>(*this).first(n);
}
// An owned series copy at an adjusted logical length: the result of with_len.
// Carries a reference to a source cache when it still serves the copied
// coefficients (a zero tail doesn't change the transform), so the source
// must outlive the result.
template <fft::engine E>
struct resized {
using T = typename E::value_type;
using engine_t = E;
static constexpr bool exact_v = false;
trunc<E> s;
std::optional<std::reference_wrapper<fft::transformed<E>>> f;
int len() const { return s.len(); }
const T& operator[](int i) const { return s[size_t(i)]; }
operator std::span<const T>() const { return std::span<const T>(s); }
operator span<E, false>() const { return s; }
maybe_cached<E, false> first(int n) const {
span<E, false> v = s;
if (n == len() && f) return {v, f->get()};
return maybe_cached<E, false>(v.first(n));
}
std::optional<std::reference_wrapper<fft::transformed<E>>> cache_opt() const { return f; }
};
// copy any operand to logical length n: extending zero-fills, shrinking truncates
template <like S>
resized<typename S::engine_t> with_len(const S& s, int n) {
using E = typename S::engine_t;
using T = typename E::value_type;
auto p = s.first(std::min(n, s.len()));
resized<E> r;
r.s.assign(size_t(n), T{});
std::span<const T> pc(p);
std::copy(pc.begin(), pc.end(), r.s.begin());
r.f = detail::cache_of(p);
return r;
}
// carries memoized transforms of power-of-two prefixes (see prefix_cached):
// product operands truncate to a covered scale to reuse them.
// Trunc-only: an exact operand participates whole, so has_cache covers it.
template <typename S>
concept has_prefix_cache = like<S> && !S::exact_v && requires(const S& s, int n) {
{ s.prefix_cache(n) } -> std::same_as<fft::transformed<typename S::engine_t>&>;
};
// Wrapper around vec which caches the transform of the whole series.
// Ops exploit the cache whenever the whole span participates; a trunc series'
// whole-sequence transform is still useful for middle products and repeated
// full-precision use.
template <fft::engine E, bool exact_>
struct cached {
using T = typename E::value_type;
using engine_t = E;
static constexpr bool exact_v = exact_;
cached() = default;
// moving coefficients in or out is free: implicit on rvalues, explicit copy otherwise
cached(vec<E, exact_>&& s_) : s(std::move(s_)) {}
explicit cached(const vec<E, exact_>& s_) : s(s_) {}
operator vec<E, exact_>() && { return std::move(s); }
int len() const { return s.len(); }
// unwrap to the owned coefficients
const vec<E, exact_>& uncached() const { return s; }
const T& operator[](int i) const { return s[size_t(i)]; }
auto begin() const { return s.cbegin(); }
auto end() const { return s.cend(); }
operator span<E, exact_>() const { return s; }
maybe_cached<E, exact_> first(int n) const {
return n == len() ? maybe_cached<E, exact_>(s, f) : maybe_cached<E, exact_>(s.first(n));
}
// the transform of the coefficients, fed to the cached fft:: entry points alongside them
fft::transformed<E>& cache() const { return f; }
template <like S>
friend bool operator==(const cached& a, const S& b) {
span<E, S::exact_v> bs = b;
return a.len() == bs.len() && std::equal(a.s.begin(), a.s.end(), bs.begin());
}
private:
vec<E, exact_> s;
mutable fft::transformed<E> f; // memoized transform: filling it is logically const
};
template <fft::engine E> using cached_exact = cached<E, true>;
template <fft::engine E> using cached_trunc = cached<E, false>;
namespace detail {
// Normalize a whole-span operand to a cached_span: the coefficients borrowed
// together with the cache serving them (the operand's own, or tmp otherwise).
// The whole-span multiply/square/middle_product paths run entirely on this form.
template <like S>
cached_span<typename S::engine_t, S::exact_v> as_cached_span(const S& s, fft::transformed<typename S::engine_t>& tmp) {
auto co = cache_of(s);
return {s, co ? co->get() : tmp};
}
/* namespace detail */ }
// Newton inversion: 1/a mod x^a.len(). Generic over any engine; per doubling step
// n -> m = 2n this is 5 transforms of size m, reusing b's transform for both circular
// products; in each product the wraparound only contaminates coefficients [0, n)
// which are already known.
//
// This is correct for non-commutative rings.
// TODO: reuse/populate the operand's whole/prefix transform caches
template <trunc_like S>
trunc<typename S::engine_t> ps_inv(const S& a) {
using E = typename S::engine_t;
using T = typename E::value_type;
int N = a.len();
trunc<E> r(size_t(N), T{});
if (N == 0) return r;
int s = nextPow2(N);
std::vector<T> b(size_t(s), T{});
b[0] = inv(a[0]);
for (int n = 1; n < N; n *= 2) {
int m = 2 * n;
auto ta = E::transform(a.first(std::min(N, m)), m);
auto tb = E::transform(std::span<const T>(b).first(n), m);
// e = a*b mod x^m; only e[n..m) is needed (and is wraparound-free).
auto e = fft::buffer_pool<T>::get(m);
E::finish(E::mul(ta, tb, m), e.span());
for (int i = 0; i < n; i++) e[i] = T{};
auto te = E::transform(std::span<const T>(e.span()), m);
auto c = fft::buffer_pool<T>::get(m);
// b' = 2b - b*(a*b): keep b on the left of e = a*b
E::finish(E::mul(tb, te, m), c.span());
for (int i = n; i < std::min(m, N); i++) b[i] = -c[i];
}
std::copy(b.begin(), b.begin() + N, r.begin());
return r;
}
// TODO: operator / can be done slightly faster than ps_inv:
// we only need the n/2 terms of ps_inv(), and can do the last Newton step directly on the quotient
// Both consume whole-sequence transforms by nature (the full span always
// participates), so only whole caches apply, never prefix caches.
template <like A>
auto square(const A& a) {
using E = typename A::engine_t;
using T = typename E::value_type;
fft::transformed<E> ta_;
auto av = detail::as_cached_span(a, ta_);
if constexpr (A::exact_v) {
// like operator*, an exact square returns has_cache, adopting the
// pointwise product as the result's transform when the engine supports it
std::vector<T> coeffs;
fft::transformed<E> f;
fft::square_cached<E>(av, av.cache(), coeffs, f);
cached<E, true> w(exact<E>(std::move(coeffs)));
w.cache() = std::move(f);
return w;
} else {
trunc<E> r(size_t(a.len()), T{});
fft::square<E>(av, av.cache(), std::span<T>(r));
return r;
}
}
// a*b + c*d, all exact; returns has_cache, adopting the summed pointwise
// product as the result's transform when the engine supports it. Reuses each
// operand's whole cache. Requires a*b and c*d to have equal length.
template <like A, like B, like C, like D>
requires fft::same_engine<A, B> && fft::same_engine<A, C> && fft::same_engine<A, D>
&& A::exact_v && B::exact_v && C::exact_v && D::exact_v
cached<typename A::engine_t, true> multiply_add2(
const A& a, const B& b, const C& c, const D& d) {
using E = typename A::engine_t;
using T = typename E::value_type;
fft::transformed<E> ta_, tb_, tc_, td_;
auto av = detail::as_cached_span(a, ta_), bv = detail::as_cached_span(b, tb_);
auto cv = detail::as_cached_span(c, tc_), dv = detail::as_cached_span(d, td_);
std::vector<T> coeffs;
fft::transformed<E> f;
fft::multiply_add2_cached<E>(
av, av.cache(),
bv, bv.cache(),
cv, cv.cache(),
dv, dv.cache(),
coeffs, f
);
cached<E, true> w(exact<E>(std::move(coeffs)));
w.cache() = std::move(f);
return w;
}
// coefficients [b.len()-1, a.len()) of a*b; requires a.len() >= b.len() > 0.
// The kernel b participates whole, so it must be exact; the result mirrors a's kind.
template <like A, exact_like B> requires fft::same_engine<A, B>
vec<typename A::engine_t, A::exact_v> middle_product(const A& a, const B& b) {
using E = typename A::engine_t;
fft::transformed<E> ta_, tb_;
auto av = detail::as_cached_span(a, ta_);
auto bv = detail::as_cached_span(b, tb_);
return vec<E, A::exact_v>(fft::middle_product<E>(
av, av.cache(),
bv, bv.cache()
));
}
namespace detail {
template <bool ea, bool eb> int product_prec(int la, int lb) {
if constexpr (ea && eb) return la > 0 && lb > 0 ? la + lb - 1 : 0;
else return ea ? lb : eb ? la : std::min(la, lb);
}
// Normalize a product operand at the given precision to a borrowed series + the
// whole cache serving it: a prefix cache at scale nextPow2(prec), or the whole
// span with the operand's own cache, or a truncated span with the caller's
// throwaway cache.
// A whole cache of an over-length operand (len > prec, which pins the other,
// necessarily trunc, operand's span at exactly prec) is only worth using when
// the untruncated span doesn't grow the transform size: a 2x'd inverse
// transform costs more than the saved forward transform.
template <like S>
auto product_operand(const S& s, int prec, fft::transformed<typename S::engine_t>& tmp) {
using E = typename S::engine_t;
if constexpr (has_prefix_cache<S>) {
return s.first(std::min(s.len(), nextPow2(prec)));
} else {
span<E, S::exact_v> v = s;
int used = std::min(v.len(), prec);
if (auto co = cache_of(s)) {
if (s.len() <= prec || fft::detail::conv_size_for(s.len() + prec - 1).n
== fft::detail::conv_size_for(2 * prec - 1).n) {
return cached_span<E, S::exact_v>{v, co->get()};
}
}
return cached_span<E, S::exact_v>{v.first(used), tmp};
}
}
/* namespace detail */ }
template <like A, like B> requires fft::same_engine<A, B>
vec<typename A::engine_t, A::exact_v && B::exact_v> operator + (const A& a, const B& b) {
using T = typename A::engine_t::value_type;
int n = (A::exact_v && B::exact_v) ? std::max(a.len(), b.len())
: A::exact_v ? b.len() : B::exact_v ? a.len() : std::min(a.len(), b.len());
vec<typename A::engine_t, A::exact_v && B::exact_v> r(size_t(n), T(0));
for (int i = 0; i < n; i++) {
r[i] = (i < a.len() ? a[i] : T(0)) + (i < b.len() ? b[i] : T(0));
}
return r;
}
template <like A, like B> requires fft::same_engine<A, B>
vec<typename A::engine_t, A::exact_v && B::exact_v> operator - (const A& a, const B& b) {
using T = typename A::engine_t::value_type;
int n = (A::exact_v && B::exact_v) ? std::max(a.len(), b.len())
: A::exact_v ? b.len() : B::exact_v ? a.len() : std::min(a.len(), b.len());
vec<typename A::engine_t, A::exact_v && B::exact_v> r(size_t(n), T(0));
for (int i = 0; i < n; i++) {
r[i] = (i < a.len() ? a[i] : T(0)) - (i < b.len() ? b[i] : T(0));
}
return r;
}
// The single multiplication operator: each operand is normalized to a borrowed
// series + whole cache (see detail::product_operand), then multiplied once.
// An exact x exact product returns a has_cache result, going through
// fft::multiply_cached so the pointwise product is adopted as the result's
// transform whenever the engine supports it.
template <like A, like B> requires fft::same_engine<A, B>
auto operator * (const A& a, const B& b) {
using E = typename A::engine_t;
using T = typename E::value_type;
constexpr bool ea = A::exact_v, eb = B::exact_v;
int prec = detail::product_prec<ea, eb>(a.len(), b.len());
if (prec == 0 || a.len() == 0 || b.len() == 0) {
if constexpr (ea && eb) return cached<E, true>{};
else return trunc<E>(size_t(prec), T(0));
}
fft::transformed<E> ta_, tb_;
auto va = detail::product_operand(a, prec, ta_);
auto vb = detail::product_operand(b, prec, tb_);
if constexpr (ea && eb) {
std::vector<T> coeffs;
fft::transformed<E> f;
auto ca = detail::as_cached_span(va, ta_), cb = detail::as_cached_span(vb, tb_);
fft::multiply_cached<E>(
ca, ca.cache(),
cb, cb.cache(),
coeffs, f
);
cached<E, true> w(exact<E>(std::move(coeffs)));
w.cache() = std::move(f);
return w;
} else {
trunc<E> r(size_t(prec), T(0));
auto ca = detail::as_cached_span(va, ta_), cb = detail::as_cached_span(vb, tb_);
fft::multiply<E>(
ca, ca.cache(),
cb, cb.cache(),
std::span<T>(r)
);
return r;
}
}
// Wrapper around trunc which caches transform(s[:2^k]) for all k,
// matching the doubling shape of ps_inv/exp so they can populate the caches.
// TODO: make ps_inv/exp populate these
template <fft::engine E>
struct prefix_cached {
using T = typename E::value_type;
using engine_t = E;
static constexpr bool exact_v = false;
prefix_cached() = default;
// moving coefficients in or out is free: implicit on rvalues, explicit copy otherwise
prefix_cached(trunc<E>&& s_) : s(std::move(s_)) {}
explicit prefix_cached(const trunc<E>& s_) : s(s_) {}
operator trunc<E>() && { return std::move(s); }
int len() const { return s.len(); }
// unwrap to the owned coefficients
const trunc<E>& uncached() const { return s; }
const T& operator[](int i) const { return s[size_t(i)]; }
auto begin() const { return s.cbegin(); }
auto end() const { return s.cend(); }
operator span<E, false>() const { return s; }
// extend precision: appends coefficients, keeping all covering caches valid
void append(std::span<const T> tail) {
s.insert(s.end(), tail.begin(), tail.end());
}
// The first k coefficients, with the covering prefix cache when one lines
// up with k (k a power of two, or k == len()); uncached otherwise.
maybe_cached<E, false> first(int k) const {
assert(k <= len());
span<E, false> v = s.first(k);
int n = nextPow2(k);
if (std::min(n, len()) == k) return {v, prefix_cache(n)};
return maybe_cached<E, false>(v);
}
// the whole-sequence transform: the prefix cache covering all of len()
fft::transformed<E>& cache() const { return prefix_cache(nextPow2(len())); }
// cache over the prefix of length min(n, len()); n a power of two
fft::transformed<E>& prefix_cache(int n) const {
assert(n > 0 && !(n & (n-1)));
int k = __builtin_ctz(unsigned(n));
if (k >= sz(caches)) caches.resize(size_t(k) + 1);
auto& c = caches[k];
int e = std::min(n, len());
if (c.len != e) {
c.t = E::transform(s.first(e), 2 * n);
c.len = e;
}
return c.t;
}
private:
trunc<E> s;
// memoized transforms: logically const; len tracks how much of s each covers
struct entry { fft::transformed<E> t; int len = 0; };
mutable std::vector<entry> caches;
};
/* namespace ecnerwala::series */ }
#line 14 "src/fft/poly.hpp"
namespace ecnerwala::poly {
// polynomial class
// As above, we represent polynomials by a series::exact containing the coefficients in reverse order.
// This representation should be internal-only:
// all accesses/constructors use the logical order though: P[k] = [x^k] P.
// To use the representation, use rev_series() / from_rev_series()
template <fft::engine E> struct vec {
using T = typename E::value_type;
using engine_t = E;
series::exact<E> c;
vec() = default;
// zero polynomial with `len` coefficient slots
explicit vec(int len) : c(size_t(len), T{}) {}
// coefficient (x^0-first) order
vec(std::initializer_list<T> coeffs) : c(std::rbegin(coeffs), std::rend(coeffs)) {}
explicit vec(std::span<const T> coeffs) : c(coeffs.rbegin(), coeffs.rend()) {}
const series::exact<E>& rev_series() const { return c; }
static vec from_rev_series(series::exact<E> s) {
vec r;
r.c = std::move(s);
return r;
}
// This should rarely be used
series::trunc<E> unrev_series(int n) const {
series::trunc<E> r(size_t(n), T{});
std::copy(begin(), begin() + std::min(n, len()), r.begin());
return r;
}
// logical (coefficient) order
auto begin() { return c.rbegin(); }
auto end() { return c.rend(); }
auto begin() const { return c.rbegin(); }
auto end() const { return c.rend(); }
int len() const { return c.len(); }
int degree() const { return len() - 1; }
T& operator[](int i) { return c[len() - 1 - i]; }
const T& operator[](int i) const { return c[len() - 1 - i]; }
T leading() const { return c.front(); }
// multiply by x^k: appends the new zero constant terms to the storage
void shift(int k = 1) {
if (len() > 0) c.insert(c.end(), size_t(k), T(0));
}
// grow (zero-filled leading coefficients) or shrink to n coefficients
void resize(int n) {
if (n >= len()) c.insert(c.begin(), size_t(n - len()), T(0));
else c.erase(c.begin(), c.begin() + (len() - n));
}
T operator()(const T& x) const {
T r{};
for (const T& v : c) r = r * x + v;
return r;
}
vec& operator+=(const vec& o) {
if (o.len() > len()) resize(o.len());
for (int i = 0; i < o.len(); i++) (*this)[i] += o[i];
return *this;
}
friend vec operator+(vec a, const vec& b) { a += b; return a; }
vec& operator-=(const vec& o) {
if (o.len() > len()) resize(o.len());
for (int i = 0; i < o.len(); i++) (*this)[i] -= o[i];
return *this;
}
friend vec operator-(vec a, const vec& b) { a -= b; return a; }
friend bool operator==(const vec& a, const vec& b) { return a.c == b.c; }
vec& operator*=(const T& n) { for (T& v : c) v *= n; return *this; }
friend vec operator*(vec a, const T& n) { a *= n; return a; }
friend vec operator*(const T& n, vec a) { a *= n; return a; }
vec& operator*=(const vec& o) { return *this = (*this) * o; }
};
// any polynomial representation exposing its reversed coefficient series
template <typename P>
concept like = requires(const P& p) {
typename P::engine_t;
{ p.len() } -> std::same_as<int>;
p.rev_series();
requires series::like<std::remove_cvref_t<decltype(p.rev_series())>>;
requires std::remove_cvref_t<decltype(p.rev_series())>::exact_v;
};
// immutable polynomial carrying the whole-sequence transform of its rev_series
template <fft::engine E>
struct cached {
using T = typename E::value_type;
using engine_t = E;
cached() = default;
// moving coefficients in or out is free: implicit on rvalues, explicit copy otherwise
cached(vec<E>&& p) : c(std::move(p.c)) {}
explicit cached(const vec<E>& p) : c(p.c) {}
operator vec<E>() && { return vec<E>::from_rev_series(std::move(c)); }
const series::cached_exact<E>& rev_series() const { return c; }
static cached from_rev_series(series::cached_exact<E> s) {
cached r;
r.c = std::move(s);
return r;
}
int len() const { return c.len(); }
int degree() const { return len() - 1; }
const T& operator[](int i) const { return c[len() - 1 - i]; }
T leading() const { return c[0]; }
T operator()(const T& x) const {
T r{};
for (const T& v : c) r = r * x + v;
return r;
}
private:
series::cached_exact<E> c;
};
// rev(a*b) = rev(a)*rev(b); the series product reuses/adopts transforms
template <like A, like B> requires fft::same_engine<A, B>
cached<typename A::engine_t> operator*(const A& a, const B& b) {
return cached<typename A::engine_t>::from_rev_series(a.rev_series() * b.rev_series());
}
template <like A>
cached<typename A::engine_t> square(const A& a) {
return cached<typename A::engine_t>::from_rev_series(square(a.rev_series()));
}
// rev(a*b + c*d) = rev(a)*rev(b) + rev(c)*rev(d)
template <like A, like B, like C, like D>
requires fft::same_engine<A, B> && fft::same_engine<A, C> && fft::same_engine<A, D>
cached<typename A::engine_t> multiply_add2(
const A& a, const B& b, const C& c, const D& d) {
return cached<typename A::engine_t>::from_rev_series(
multiply_add2(a.rev_series(), b.rev_series(), c.rev_series(), d.rev_series()));
}
template <like A, like B> requires fft::same_engine<A, B>
bool operator==(const A& a, const B& b) {
return a.rev_series() == b.rev_series();
}
// finite-support linear form
// These are one side of the pairing <vec P, series::vec S> = [x^0] P(1/x) S(x).
// (Strictly speaking, this is actually <>_d where we take polynomials of degree < d.)
// The main point of this wrapper is that if we have <*, S> and want <P *, S>, that's a middle product by P.
//
// TODO: Should we split it apart into <*, S> and <P, *>?
//
// Some use cases of this pairing:
// <P, 1/(1-ax)> = P(a)
// if we represent P as a "polynomial" in the differential operator D (x^k = k! D^k):
// <P, e^{aD}> = P(a)
template <fft::engine E>
struct form {
using T = typename E::value_type;
// coeffs of S in <*, S>; always whole-cached: the kernel transform is
// what repeated middle products against the same form reuse
series::cached_exact<E> c;
form() = default;
explicit form(int len) : c(series::exact<E>(size_t(len), T{})) {}
// We don't provide coefficient-list constructors, to avoid ordering confusion.
const series::cached_exact<E>& rev_series() const { return c; }
static form from_rev_series(series::cached_exact<E> s) {
form r;
r.c = std::move(s);
return r;
}
static form from_poly(const vec<E>& p) { return from_rev_series(series::cached_exact<E>(p.rev_series())); }
int len() const { return c.len(); }
// Restrict the form's domain: only valid against exact series of length n
form for_length(int n) const {
auto r = series::exact<E>(c);
if (n >= len()) r.insert(r.begin(), size_t(n - len()), T(0));
else r.erase(r.begin(), r.begin() + (len() - n));
return from_rev_series(std::move(r));
}
// the functional p -> p(z) on polynomials of length up to len (weight z^i on [x^i])
static form polynomial_evaluation(T z, int len) {
series::exact<E> k(size_t(len), T{});
T p = T(1);
for (int i = 0; i < len; i++) { k[i] = p; p *= z; }
return from_rev_series(std::move(k));
}
template <like P>
T operator()(const P& p) const {
assert(p.len() <= len());
T r{};
for (int i = 0; i < p.len(); i++) r += c[i] * p[i]; // weights multiply from the left
return r;
}
// <*, S> -> <q x *, S>
template <like P>
form composed_with(const P& q) const {
assert(q.len() > 0 && q.len() <= len());
return from_rev_series(middle_product(c, q.rev_series()));
}
// <P, *> -> <P, s x *>
template <series::like S> requires std::same_as<typename S::engine_t, E>
form composed_with(const S& s) const {
if constexpr (!S::exact_v) assert(s.len() >= len());
series::vec<E, S::exact_v> r = c * s;
r.resize(size_t(len()));
return from_rev_series(series::exact<E>(std::move(r)));
}
};
// ==== multipoint evaluation / interpolation ====
// Subproduct tree over points a[0:N]
// BFS-order tree, each node holds prod (x - a[i]) as a cached vec.
template <fft::engine E>
struct subproduct_tree {
using T = typename E::value_type;
int N;
std::vector<cached<E>> nodes;
explicit subproduct_tree(std::span<const T> pts) : N(sz(pts)), nodes(size_t(2) * N) {
assert(N > 0);
for (int i = 0; i < N; i++) {
nodes[N + i] = vec<E>{-pts[i], T(1)};
}
for (int i = N - 1; i > 0; i--) {
nodes[i] = nodes[2*i] * nodes[2*i+1];
}
}
// number of points under node i
int size(int i) const { return nodes[i].len() - 1; }
// prod (x - z_j) over node i's leaves; length size(i) + 1
const cached<E>& prod(int i) const { return nodes[i]; }
// Computes, for each i, f(product_{j != i} (1 - a[j] x)). Requires f.len() == N.
std::vector<T> pushdown(form<E> f) const {
assert(f.len() == N);
std::vector<form<E>> down(size_t(2) * N);
down[1] = std::move(f);
for (int i = 1; i < N; i++) {
// the form's kernel transform serves both children's middle products
down[2*i+0] = down[i].composed_with(nodes[2*i+1]);
down[2*i+1] = down[i].composed_with(nodes[2*i+0]);
down[i] = form<E>{}; // done with the parent; free it early
}
std::vector<T> out(size_t(N), T{});
for (int i = 0; i < N; i++) out[i] = down[N + i].rev_series()[0];
return out;
}
// Compute sum_i leaf_vals[i] prod_{j!=i} (x - a[j]) (transpose of pushdown)
cached<E> combine_up(std::span<const T> leaf_vals) const {
assert(sz(leaf_vals) == N);
std::vector<cached<E>> up(size_t(2) * N);
for (int i = 0; i < N; i++) {
up[N + i] = vec<E>{leaf_vals[i]};
}
for (int i = N - 1; i > 0; i--) {
up[i] = multiply_add2(up[2*i+0], nodes[2*i+1], up[2*i+1], nodes[2*i+0]);
up[2*i+0] = cached<E>{};
up[2*i+1] = cached<E>{};
}
return std::move(up[1]);
}
};
template <fft::engine E>
std::vector<typename E::value_type> multipoint(
const vec<E>& p,
std::span<const typename E::value_type> pts
) {
if (pts.empty()) return {};
int N = sz(pts);
subproduct_tree<E> tree{pts};
auto q = series::trunc<E>(tree.prod(1).rev_series());
q.resize(p.len()); // inverse precision must cover the form's window
form<E> f = form<E>::from_poly(p).composed_with(ps_inv(q));
return tree.pushdown(f.for_length(N));
}
template <fft::engine E>
vec<E> interpolate(
std::span<const typename E::value_type> pts,
std::span<const typename E::value_type> vals
) {
using T = typename E::value_type;
assert(sz(pts) == sz(vals));
if (pts.empty()) return {};
int N = sz(pts);
using ps = series::trunc<E>;
subproduct_tree<E> tree{pts};
auto root = ps(tree.prod(1).rev_series());
root.shrink(N);
// We need to evaluate the derivative of the root at each point
ps deriv_root = root;
for (int i = 0; i < N; i++) {
deriv_root[i] *= T(N - i);
}
std::vector<T> denoms = tree.pushdown(
form<E>::from_rev_series(series::exact<E>(ps_inv(root) * deriv_root))
);
std::vector<T> leaf_vals(size_t(N), T{});
for (int i = 0; i < N; i++) leaf_vals[i] = vals[i] / denoms[i];
return tree.combine_up(std::span<const T>(leaf_vals));
}
/* namespace ecnerwala::poly */ }
#line 2 "src/fft/series.hpp"
#line 9 "src/fft/series.hpp"
#line 11 "src/fft/series.hpp"
// ==== analytic ops ====
// Free functions over series-like operands; each borrows the operand's span
// and writes a fresh result.
// TODO: reuse/populate the operands' whole/prefix transform caches
namespace ecnerwala::series {
template <like S>
vec<typename S::engine_t, S::exact_v> stretch(const S& a_, int n) {
using E = typename S::engine_t;
span<E, S::exact_v> a = a_;
vec<E, S::exact_v> r(size_t(a.len()));
for (int i = 0; i*n < a.len(); i++) {
r[i*n] = a[i];
}
return r;
}
template <like S>
vec<typename S::engine_t, S::exact_v> deriv_shift(const S& a_) {
using E = typename S::engine_t;
span<E, S::exact_v> a = a_;
vec<E, S::exact_v> r(a.begin(), a.end());
for (int i = 0; i < r.len(); i++) {
r[i] *= i;
}
return r;
}
template <like S>
vec<typename S::engine_t, S::exact_v> integ_shift(const S& a_) {
using E = typename S::engine_t;
using T = typename E::value_type;
span<E, S::exact_v> a = a_;
assert(a[0] == 0);
vec<E, S::exact_v> r(a.begin(), a.end());
T f = 1;
for (int i = 1; i < r.len(); i++) {
r[i] *= f;
f *= i;
}
f = inv(f);
for (int i = r.len() - 1; i > 0; i--) {
r[i] *= f;
f *= i;
}
return r;
}
template <like S>
vec<typename S::engine_t, S::exact_v> integ_shift_offset(const S& a_, int offset) {
using E = typename S::engine_t;
using T = typename E::value_type;
span<E, S::exact_v> a = a_;
vec<E, S::exact_v> r(a.begin(), a.end());
T f = 1;
for (int i = 0; i < r.len(); i++) {
r[i] *= f;
f *= i + offset;
}
assert(f != 0);
f = inv(f);
for (int i = r.len() - 1; i >= 0; i--) {
r[i] *= f;
f *= i + offset;
}
return r;
}
template <trunc_like S>
trunc<typename S::engine_t> deriv_shift_log(const S& a) {
return deriv_shift(a) * ps_inv(a);
}
template <trunc_like S>
trunc<typename S::engine_t> ps_log(const S& a) {
assert(a[0] == 1);
return integ_shift(deriv_shift_log(a));
}
template <trunc_like S>
trunc<typename S::engine_t> ps_exp(const S& a_) {
// See https://mathexp.eu/bostan/publications/BoSc09a.pdf for details
using E = typename S::engine_t;
using T = typename E::value_type;
span<E, false> a = a_;
assert(a.len() >= 1);
assert(a[0] == 0);
trunc<E> r(1, T(1)); r.reserve(size_t(a.len()));
trunc<E> invR(1, T(1)); invR.reserve(size_t(a.len()));
while (r.len() < a.len()) {
int o_sz = r.len();
int n_sz = std::min(o_sz * 2, a.len());
trunc<E> t = deriv_shift(trunc<E>(a.begin(), a.begin() + o_sz));
fft::multiply_circular<E>(std::span<const T>(t), std::span<const T>(r).first(o_sz), std::span<T>(t), o_sz);
t = deriv_shift(r) - t;
t *= invR;
t.resize(size_t(n_sz - o_sz));
trunc<E> v(a.begin() + o_sz, a.begin() + n_sz);
v -= integ_shift_offset(t, o_sz);
v *= r;
r.resize(size_t(n_sz));
std::copy(v.begin(), v.end(), r.begin() + o_sz);
if (r.len() < a.len()) {
// double invR via a Newton step
assert(r.len() == 2 * invR.len());
int n = invR.len();
int nn = r.len();
trunc<E> tmp(size_t(4) * n);
fft::square<E>(std::span<const T>(invR).first(n), std::span<T>(tmp));
fft::multiply<E>(std::span<const T>(tmp).first(nn), std::span<const T>(r).first(nn), std::span<T>(tmp));
invR.resize(size_t(nn));
for (int i = n; i < nn; i++) invR[i] = -tmp[i];
}
}
return r;
}
template <trunc_like S>
trunc<typename S::engine_t> ps_pow_monic(const S& a_, typename S::engine_t::value_type k) {
using E = typename S::engine_t;
span<E, false> a = a_;
if (a.len() == 0) return {};
assert(a[0] == 1);
trunc<E> l = ps_log(a_);
l *= k;
return ps_exp(l);
}
template <trunc_like S>
trunc<typename S::engine_t> ps_pow(const S& a_, int64_t k) {
using E = typename S::engine_t;
using T = typename E::value_type;
span<E, false> a = a_;
assert(k >= 0);
if (k == 0) {
trunc<E> r(size_t(a.len()), T(0));
if (r.len() > 0) r[0] = T(1);
return r;
}
int st = 0;
while (st < a.len() && a[st] == 0) st++;
if (st > 0 && k > (a.len() - 1) / st) {
return trunc<E>(size_t(a.len()), T(0));
}
trunc<E> r(a.begin() + st, a.end() - (st * (k-1)));
T leading_coeff = r[0];
r *= inv(leading_coeff);
r = ps_pow_monic(r, T(k));
r *= power(leading_coeff, k);
r.insert(r.begin(), size_t(st * k), T(0));
assert(r.len() == a.len());
return r;
}
template <trunc_like S>
trunc<typename S::engine_t> to_newton_sums(const S& a, int deg) {
auto r = deriv_shift_log(a);
r[0] = deg;
for (int i = 1; i < r.len(); i++) r[i] = -r[i];
return r;
}
template <trunc_like S>
trunc<typename S::engine_t> from_newton_sums(const S& s_, int deg) {
using E = typename S::engine_t;
span<E, false> s = s_;
assert(s[0] == deg);
trunc<E> r(s.begin(), s.end());
r[0] = 0;
for (int i = 1; i < r.len(); i++) r[i] = -r[i];
return ps_exp(integ_shift(std::move(r)));
}
// Calculates prod 1/(1-x^i)^{a[i]}
template <trunc_like S>
trunc<typename S::engine_t> euler_transform(const S& a) {
using E = typename S::engine_t;
trunc<E> r = deriv_shift(a);
std::vector<bool> is_prime(size_t(r.len()), true);
for (int p = 2; p < r.len(); p++) {
if (!is_prime[p]) continue;
for (int i = 1; i*p < r.len(); i++) {
r[i*p] += r[i];
is_prime[i*p] = false;
}
}
return ps_exp(integ_shift(r));
}
template <trunc_like S>
trunc<typename S::engine_t> inverse_euler_transform(const S& a) {
using E = typename S::engine_t;
trunc<E> r = deriv_shift(ps_log(a));
std::vector<bool> is_prime(size_t(r.len()), true);
for (int p = 2; p < r.len(); p++) {
if (!is_prime[p]) continue;
for (int i = (r.len()-1)/p; i >= 1; i--) {
r[i*p] -= r[i];
is_prime[i*p] = false;
}
}
return integ_shift(r);
}
// Helper packed bivariate buffer for Kinoshita-Li composition (arXiv:2404.05177).
//
// The motivation is performing Bostan-Mori (Graeffe root-squaring) to compute
// something like [x^n] P / Q_0(x, y) with deg_y(Q_0) = 1 and deg_x(Q_0) = n.
//
// In each step, we want to compute Q_{i+1}(x^2, y) = Q_i(x, y) * Q_i(-x, y).
// This doubles the degree of y and also lets us truncate x at half the previous
// degree, leaving the total size invariant.
//
// We will store Q as a packed buffer with x as the inner dimension to facilitate easy Q(-x) substitution.
// The inner span will be 2*deg(x), and the outer span will be 2*deg(y).
// As we advance, we will also return the cached transform of Q_i(-x, y) for the caller to use in the numerator.
template <fft::engine E> struct packed_bivariate {
using T = typename E::value_type;
int L, l;
std::vector<T> c;
// Q_0 = 1 - y g(x), deg g < n <= 2^L
packed_bivariate(int L_, std::span<const T> g) : L(L_), l(0), c(size_t(4) << L) {
c[0] = T(1);
for (int i = 0; i < sz(g); i++) c[(2 << L) + i] = -g[i];
}
fft::transformed<E> advance() {
int B = 4 << L;
auto tq = E::transform(std::span<const T>(c), B);
auto tn = E::negate_arg(tq, B);
E::finish(
E::downsample(E::mul(tq, tn, B), B/2, false),
std::span<T>(c).first(B/2)
);
l++;
// undo the circular wraparound using monicity in y
for (int i = 0; i < (2 << (L - l)); i++) {
c[(2 << L) + i] = c[i];
c[i] = T(0);
}
c[2 << L] -= T(1);
c[0] = T(1);
// zero x coefficients beyond the level's truncation mod x^(2^(L-l))
std::fill(c.begin() + (2 << L) + (1 << (L - l)), c.end(), T(0));
for (int i = 0; i < (2 << L); i += 2 << (L - l)) {
for (int j = 0; j < (1 << (L - l)); j++) {
c[i + (1 << (L - l)) + j] = T(0);
}
}
return tn;
}
};
// Calculates f(g(x)) mod x^n where deg(g) == n
template <trunc_like SF, trunc_like SG> requires fft::same_engine<SF, SG>
trunc<typename SF::engine_t> ps_compose(const SF& f_, const SG& g_) {
using E = typename SF::engine_t;
using T = typename E::value_type;
span<E, false> f = f_;
span<E, false> g = g_;
if (g.len() == 0) return {};
int m = f.len();
int n = g.len();
// https://arxiv.org/pdf/2404.05177
// Consider P(y) = f(1/y) has terms from y^{-(m-1)}...y^0 (Laurent series)
// We want [y^0] P(y) / (1 - y g(x))
// Let Q_0 = 1 - yg(x)
// Q_{i+1}(x^2, y) = Q_i(x, y) * Q_i(-x, y) mod x^{ceil(n / 2^i)}
// deg_y(Q_i) = 2^i, deg_x(Q_i) = ceil(n / 2^i) - 1
//
// [y^0] P(y) / Q_l(x^2^l, y) * Q_{l-1}(-x^2^{l-1}, y) * Q_{l-2}(-x^2^{l-2}, y) * ... * Q_0(-x, y)
// The total y deg of Q_{k-1} ... Q_0 is 2^k-1
int L = __builtin_ctz(unsigned(nextPow2(n)));
int B = 4 << L;
packed_bivariate<E> Q(L, g.coeffs());
// tneg[l] is the transform of Q_l(-x, y), reused by the pushdown pass below
std::vector<fft::transformed<E>> tneg;
tneg.reserve(L);
for (int l = 1; l <= L; l++) tneg.push_back(Q.advance());
trunc<E> P;
{
P = trunc<E>(f.begin(), f.end());
std::reverse(P.begin(), P.end());
trunc<E> QL((1 << L) + 1);
for (int i = 0; i <= (1 << L); i++) {
QL[i] = Q.c[2 * i];
}
QL.resize(size_t(m), T(0));
P *= ps_inv(QL);
std::reverse(P.begin(), P.end());
P.resize(size_t(1) << L, T(0));
std::reverse(P.begin(), P.end());
P.resize(size_t(B), T(0));
for (int i = (1 << L) - 1; i > 0; i--) {
P[2*i] = P[i];
P[i] = T(0);
}
}
for (int l = L-1; l >= 0; l--) {
// Spread it out, clear the high terms
for (int i = (2 << L) - 1; i > 0; i--) {
T v = P[i];
P[2*i] = ((2*i) & (1 << (L-l))) ? T(0) : v;
P[i] = T(0);
}
auto tp = E::transform(std::span<const T>(P), B);
E::finish(E::mul(tneg[l], tp, B), std::span<T>(P));
for (int i = 0; i < (2 << L); i++) {
P[i] = P[(2 << L) + i];
P[(2 << L) + i] = T(0);
}
}
return trunc<E>(P.begin(), P.begin() + n);
}
// [x^k] p(x)/q(x) (Bostan-Mori) for an exact rational function.
template <exact_like P, exact_like Q> requires fft::same_engine<P, Q>
P::engine_t::value_type kth_term_of_rational_function(
const P& p,
const Q& q,
uint64_t k
) {
using E = P::engine_t;
using T = E::value_type;
assert(q.len() > 0 && q[0] != T(0));
// Check this here so we avoid accessing p[0]
if (p.len() == 0) return T(0);
// Size up in a pretty conservative way
int d = std::max(p.len() + 1, q.len());
assert(d >= 2);
int n = nextPow2((d-1) + d - 1); // >= d
// Seed the loop transforms from any whole caches; the buffers below hold the
// current p, q (zero-padded, which extend_to tolerates).
fft::transformed<E> tq, tp;
if (auto cq = detail::cache_of(q)) { E::extend_to(cq->get(), n, q); tq = cq->get(); }
if (auto cp = detail::cache_of(p)) { E::extend_to(cp->get(), n, p); tp = cp->get(); }
std::vector<T> p_buf(d-1, T(0));
std::ranges::copy(std::span<const T>(p), p_buf.begin());
std::vector<T> q_buf(d, T(0));
std::ranges::copy(std::span<const T>(q), q_buf.begin());
while (k > 0) {
E::extend_to(tq, n, q_buf);
auto tnq = E::negate_arg(tq, n);
E::extend_to(tp, n, p_buf);
// P <- downsample(P(x) * Q(-x))
auto ntp = E::downsample(E::mul(tp, tnq, n), n/2, bool(k & 1));
assert(ntp.size() == n/2);
if constexpr (std::same_as<typename E::product, typename E::transformed>) {
tp = ntp;
} else {
tp = {};
}
E::finish(std::move(ntp), std::span(p_buf));
k >>= 1;
// Save the last iteration if we're done
if (!k) {
// HACK: fix the constant coefficient of q only
q_buf[0] *= q_buf[0];
break;
}
// Q <- downsample(Q(x) * Q(-x))
auto ntq = E::downsample(E::mul(tq, tnq, n), n/2, false);
assert(ntq.size() == n/2);
if constexpr (std::same_as<typename E::product, typename E::transformed>) {
tq = ntq;
} else {
tq = {};
}
if (n/2 == d-1) {
// Fix the wraparound
T v0 = q_buf[0] * q_buf[0];
E::finish(std::move(ntq), std::span(q_buf).first(d-1));
q_buf[d-1] = std::exchange(q_buf[0], v0) - v0;
} else {
E::finish(std::move(ntq), std::span(q_buf));
}
}
return p_buf[0] * inv(q_buf[0]);
}
// Find the kth term of linearly recurrent sequence S with char poly Q and len(S) >= len(Q)-1
template <trunc_like S, exact_like Q> requires fft::same_engine<S, Q>
S::engine_t::value_type kth_term_of_linear_recurrence(
const S& s,
const Q& q,
uint64_t k
) {
using E = S::engine_t;
using T = E::value_type;
assert(q.len() > 0 && q[0] != T(0));
assert(s.len() >= q.len()-1);
// Don't even bother with P so we don't have to do truncation checks
// TODO: Could use generic multiply for this whole part?
fft::transformed<E> tq;
auto q_cached = detail::as_cached_span(q, tq);
// Compute the prefix and then hard-cast it to exact
span<E, false> sv = s;
auto p = exact<E>(sv.first(q.len()-1) * q_cached);
return kth_term_of_rational_function(p, q_cached, k);
}
/* namespace ecnerwala::series */ }
#line 2 "src/fft/test_util.test.hpp"
#line 5 "src/fft/test_util.test.hpp"
#line 2 "src/fft/engines/ntt.hpp"
#line 8 "src/fft/engines/ntt.hpp"
#line 2 "src/fft/core.hpp"
#line 8 "src/fft/core.hpp"
#line 2 "src/modnum.hpp"
#line 9 "src/modnum.hpp"
template <typename T> T mod_inv_in_range(T a, T m) {
// assert(0 <= a && a < m);
T x = a, y = m;
// abs coeff of a in x and y (they're always opposite sign)
T vx = 1, vy = 0;
bool swap = false;
while (x) {
T k = y / x;
y %= x;
vy += k * vx;
std::swap(x, y);
std::swap(vx, vy);
swap ^= 1;
}
assert(y == 1);
return swap ? vy : m - vy;
}
template <typename T> struct extended_gcd_result {
T gcd;
T coeff_a, coeff_b;
};
template <typename T> extended_gcd_result<T> extended_gcd(T a, T b) {
T x = a, y = b;
// coeff of a and b in x and y
T ax = 1, ay = 0;
T bx = 0, by = 1;
while (x) {
T k = y / x;
y %= x;
ay -= k * ax;
by -= k * bx;
std::swap(x, y);
std::swap(ax, ay);
std::swap(bx, by);
}
return {y, ay, by};
}
template <typename T> T mod_inv(T a, T m) {
a %= m;
a = a < 0 ? a + m : a;
return mod_inv_in_range(a, m);
}
// Derives the boilerplate operator surface of a number type from its compound
// ops, ==, neg(), and inv().
// Bodies are only instantiated on use, so a type may omit some of the
// underlying pieces if the corresponding derived ops are never called.
template <typename Self>
struct num_ops {
Self operator+ () const { return static_cast<const Self&>(*this); }
Self operator- () const { return static_cast<const Self&>(*this).neg(); }
friend Self operator ++ (Self& a, int) { Self r = a; ++a; return r; }
friend Self operator -- (Self& a, int) { Self r = a; --a; return r; }
friend Self operator + (const Self& a, const Self& b) { return Self(a) += b; }
friend Self operator - (const Self& a, const Self& b) { return Self(a) -= b; }
friend Self operator * (const Self& a, const Self& b) { return Self(a) *= b; }
friend Self operator / (const Self& a, const Self& b) { return Self(a) /= b; }
friend bool operator != (const Self& a, const Self& b) { return !(a == b); }
friend Self neg(const Self& a) { return a.neg(); }
friend Self inv(const Self& a) { return a.inv(); }
};
// Storage and arithmetic for numbers mod Self::MOD, as a reduced
// representative v in [0, MOD) of unsigned type V.
// The type provides static MOD (of type V), reduce (value -> representative),
// and *=;
// everything else is derived here, valid for any MOD up to V's full range
// (sums and differences are tracked mod 2^bits, so no headroom is needed).
// Hooks may be overridden in the type's own body (e.g. a faster += / -=).
template <typename Self, typename V>
struct mod_ops : num_ops<Self> {
static_assert(std::unsigned_integral<V>);
V v;
struct is_reduced_tag {};
mod_ops() : v(0) {}
mod_ops(V v_, is_reduced_tag) : v(v_) { assert(v < Self::MOD); }
template <std::integral I> mod_ops(I x) : v(Self::reduce(x)) {}
static Self from_reduced(V v) { return Self(v, is_reduced_tag{}); }
// A negative value reduces via its nonnegative complement: x = -1 - ~x.
static V reduce(std::signed_integral auto x) {
using U = std::make_unsigned_t<decltype(x)>;
return x < 0 ? V(Self::MOD - 1 - Self::reduce(U(~x))) : Self::reduce(U(x));
}
explicit operator V() const { return v; }
std::make_signed_t<V> balanced() const {
return std::make_signed_t<V>(Self::MOD-v > v ? v : v - Self::MOD);
}
friend bool operator == (const Self& a, const Self& b) { return a.v == b.v; }
friend std::ostream& operator << (std::ostream& out, const Self& n) { return out << n.v; }
friend std::istream& operator >> (std::istream& in, Self& n) { int64_t v_; in >> v_; n = Self(v_); return in; }
Self& operator ++ () {
++v;
if (v == Self::MOD) v = 0;
return self();
}
Self& operator -- () {
if (v == 0) v = Self::MOD;
--v;
return self();
}
Self& operator += (const Self& o) { v = Self::sub_mod_raw(v, Self::MOD - o.v); return self(); }
Self& operator -= (const Self& o) { v = Self::sub_mod_raw(v, o.v); return self(); }
Self& operator /= (const Self& o) { return self() *= o.inv(); }
// Returns a - b mod MOD, for b in [0, MOD]; wraparound detects the underflow.
static V sub_mod_raw(V a, V b) { return a < b ? a - b + Self::MOD : a - b; }
Self neg() const { return from_reduced(v ? Self::MOD - v : 0); }
Self inv() const { return from_reduced(mod_inv_in_range(v, Self::MOD)); }
private:
Self& self() { return static_cast<Self&>(*this); }
};
template <auto MOD_> struct modnum : mod_ops<modnum<MOD_>, std::make_unsigned_t<decltype(MOD_)>> {
using Self = modnum;
static_assert(MOD_ > 0, "MOD must be positive");
using V = std::make_unsigned_t<decltype(MOD_)>;
static constexpr V MOD = V(MOD_);
using base = mod_ops<modnum, V>;
using base::base;
using base::v;
using base::reduce;
static V reduce(std::unsigned_integral auto x) { return V(x % MOD); }
explicit operator std::make_signed_t<V>() const
requires (MOD <= V(std::numeric_limits<std::make_signed_t<V>>::max()))
{
return std::make_signed_t<V>(v);
}
Self& operator *= (const Self& o) {
if constexpr (sizeof(V) <= 4) v = V(uint64_t(v) * o.v % MOD);
else v = V(__uint128_t(v) * o.v % MOD);
return *this;
}
};
struct mod_goldilocks : mod_ops<mod_goldilocks, uint64_t> {
using Self = mod_goldilocks;
static constexpr uint64_t MOD = 0xffffffff00000001ull;
static constexpr uint64_t EPS = -MOD;
// We have 2^32 is a primitive 6th root of unity.
// Note that omega_8 + omega_8^7 == 2^24 - 2^72 == sqrt(2)
// We'll pick the root so that 2^24 - 2^72 is our primitive 384th root of unity.
static constexpr uint64_t PRIMITIVE_ROOT = 2717;
using base = mod_ops<mod_goldilocks, uint64_t>;
using base::base;
using base::reduce;
mod_goldilocks() = default;
mod_goldilocks(__int128_t a) : base(a < 0 ? uint64_t(MOD - 1 - __uint128_t(~a) % MOD) : uint64_t(__uint128_t(a) % MOD), is_reduced_tag{}) {}
mod_goldilocks(__uint128_t a) : base(uint64_t(a % MOD), is_reduced_tag{}) {}
// Avoids the division: any uint64_t is within MOD of reduced.
static uint64_t reduce(std::unsigned_integral auto x) {
static_assert(sizeof(x) <= 8);
uint64_t a = x;
return a >= MOD ? a - MOD : a;
}
// returns a-b, assuming -MOD <= a-b, e.g. b <= MOD
static uint64_t sub_mod_raw(uint64_t a, uint64_t b) {
#if defined(__x86_64__)
// TODO: We could try to write this using intrinsics, but GCC sometimes produces the wrong code.
uint64_t res_wrapped = a;
uint64_t adjustment = b;
asm (
// AT&T syntax: SRC DST
"sub %[y], %[x]\n\t"
// Trick from plonky2 implementation:
// After sub, flag CF is set iff we underflowed. We want to correct by EPS == 2^32 - 1 iff C is set.
// sbb (subtract with borrow) computes DST <- DST - SRC - CF
// Thus, we can use the 32-bit form of sbb on a dummy register to load CF ? EPS : 0.
// Here, we'll just reuse the original register holding b.
"sbb %k[y], %k[y]\n\t"
: [x] "+r"(res_wrapped),
[y] "+r"(adjustment)
:
: "cc"
);
#else
uint64_t res_wrapped = a - b;
uint64_t adjustment = (res_wrapped > a) ? EPS : 0;
#endif
return res_wrapped - adjustment;
}
// Reduce lo + 2^64 * mi + 2^96 * hi, where hi <= MOD
static uint64_t reduce_u160_raw(uint64_t lo, uint32_t mi, uint64_t hi) {
// result = lo - hi + EPS * mi
// 0 <= lo <= 2^64 - 1 = MOD + EPS - 1
// 0 <= EPS * mi <= (2^32 - 1) * EPS = MOD - 1 - EPS
// 0 <= hi <= MOD
// -MOD <= lo - hi + EPS * mi <= 2*MOD-2
// so we do have some leeway
return sub_mod_raw(sub_mod_raw(lo, hi), MOD-(uint64_t(mi)<<32)+mi);
}
static uint64_t reduce_u128_raw(__uint128_t v) {
uint64_t hi = uint64_t(v >> 64);
uint64_t lo = uint64_t(v);
uint32_t hi_hi = uint32_t(hi >> 32);
uint32_t hi_lo = uint32_t(hi);
return reduce_u160_raw(lo, hi_lo, hi_hi);
}
Self& operator *= (Self o) {
v = reduce_u128_raw(__uint128_t(v) * __uint128_t(o.v));
return *this;
}
};
template <typename T> T power(T a, long long b) {
assert(b >= 0);
T r = 1; while (b) { if (b & 1) r *= a; b >>= 1; a *= a; } return r;
}
template <typename U, typename V> struct pairnum : num_ops<pairnum<U, V>> {
using Self = pairnum;
U u;
V v;
pairnum() : u(0), v(0) {}
pairnum(long long val) : u(val), v(val) {}
pairnum(const U& u_, const V& v_) : u(u_), v(v_) {}
friend std::ostream& operator << (std::ostream& out, const Self& n) { return out << '(' << n.u << ',' << ' ' << n.v << ')'; }
friend std::istream& operator >> (std::istream& in, Self& n) { long long val; in >> val; n = Self(val); return in; }
friend bool operator == (const Self& a, const Self& b) { return a.u == b.u && a.v == b.v; }
Self inv() const {
return Self(u.inv(), v.inv());
}
Self neg() const {
return Self(u.neg(), v.neg());
}
Self& operator ++ () {
++u, ++v;
return *this;
}
Self& operator -- () {
--u, --v;
return *this;
}
Self& operator += (const Self& o) {
u += o.u;
v += o.v;
return *this;
}
Self& operator -= (const Self& o) {
u -= o.u;
v -= o.v;
return *this;
}
Self& operator *= (const Self& o) {
u *= o.u;
v *= o.v;
return *this;
}
Self& operator /= (const Self& o) {
u /= o.u;
v /= o.v;
return *this;
}
};
template <typename tag> struct dynamic_modnum : mod_ops<dynamic_modnum<tag>, uint32_t> {
using Self = dynamic_modnum;
private:
inline static uint32_t MOD_ = 0;
inline static uint64_t BARRETT_M = 0;
public:
// Make only the const-reference public, to force the use of set_mod
static constexpr uint32_t const& MOD = MOD_;
using base = mod_ops<dynamic_modnum, uint32_t>;
using base::base;
using base::v;
using base::reduce;
// Barret reduction taken from KACTL:
/**
* Author: Simon Lindholm
* Date: 2020-05-30
* License: CC0
* Source: https://en.wikipedia.org/wiki/Barrett_reduction
* Description: Compute $a \% b$ about 5 times faster than usual, where $b$ is constant but not known at compile time.
* Returns a value congruent to $a \pmod b$ in the range $[0, 2b)$.
* Status: proven correct, stress-tested
* Measured as having 4 times lower latency, and 8 times higher throughput, see stress-test.
* Details:
* More precisely, it can be proven that the result equals 0 only if $a = 0$,
* and otherwise lies in $[1, (1 + a/2^64) * b)$.
*/
static void set_mod(int mod) {
assert(mod > 0);
MOD_ = uint32_t(mod);
BARRETT_M = (uint64_t(-1) / MOD);
}
static uint32_t barrett_reduce_partial(uint64_t a) {
return uint32_t(a - uint64_t((__uint128_t(BARRETT_M) * a) >> 64) * MOD);
}
static uint32_t barrett_reduce(uint64_t a) {
int32_t res = int32_t(barrett_reduce_partial(a) - MOD);
return uint32_t((res < 0) ? res + int32_t(MOD) : res);
}
struct mod_reader {
friend std::istream& operator >> (std::istream& i, mod_reader) {
int mod; i >> mod;
Self::set_mod(mod);
return i;
}
};
static mod_reader MOD_READER() {
return mod_reader();
}
static uint32_t reduce(std::unsigned_integral auto x) {
static_assert(sizeof(x) <= 8);
return barrett_reduce(x);
}
explicit operator int() const { return int(v); }
Self& operator *= (const Self& o) {
v = barrett_reduce(uint64_t(v) * o.v);
return *this;
}
};
template <typename T> struct mod_constraint {
T v, mod;
friend mod_constraint operator & (mod_constraint a, mod_constraint b) {
if (a.mod < b.mod) std::swap(a, b);
if (b.mod == 1) return a;
extended_gcd_result<T> egcd = extended_gcd<T>(a.mod, b.mod);
assert(a.v % egcd.gcd == b.v % egcd.gcd);
T extra = b.v - a.v % b.mod;
extra /= egcd.gcd;
extra *= egcd.coeff_a;
extra %= b.mod / egcd.gcd;
extra += (extra < 0) ? b.mod / egcd.gcd : 0;
return mod_constraint{
a.v + extra * a.mod,
a.mod * (b.mod / egcd.gcd)
};
}
};
#line 11 "src/fft/core.hpp"
namespace ecnerwala::fft {
// ==== core: roots, buffers, raw transforms ====
// Complex
template <typename dbl> struct cplx { /// start-hash
dbl x, y;
cplx(dbl x_ = 0, dbl y_ = 0) : x(x_), y(y_) { }
friend cplx operator+(cplx a, cplx b) { return cplx(a.x + b.x, a.y + b.y); }
friend cplx operator-(cplx a, cplx b) { return cplx(a.x - b.x, a.y - b.y); }
friend cplx operator*(cplx a, cplx b) { return cplx(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x); }
friend cplx conj(cplx a) { return cplx(a.x, -a.y); }
friend cplx inv(cplx a) { dbl n = (a.x*a.x+a.y*a.y); return cplx(a.x/n,-a.y/n); }
};
// getRoot implementations
template <typename num> struct getRoot {
static num f(int k) = delete;
};
template <typename dbl> struct getRoot<cplx<dbl>> {
static cplx<dbl> f(int k) {
#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif
dbl a=2*M_PI/k;
return cplx<dbl>(cos(a),sin(a));
}
};
template <int MOD> struct primitive_root {
static const int value;
};
// 998244353 = (119 << 23) + 1 = 2^30 - 2^26 - 2^23 + 1
template <> struct primitive_root<998244353> {
static const int value = 3;
};
// babybear prime
template <> struct primitive_root<(15 << 27) + 1> {
static const int value = 31;
};
// koalabear prime
template <> struct primitive_root<(127 << 24) + 1> {
static const int value = 3;
};
template <> struct primitive_root<(7 << 26) + 1> {
static const int value = 3;
};
template <> struct primitive_root<(5 << 25) + 1> {
static const int value = 3;
};
template <int MOD> struct getRoot<modnum<MOD>> {
static modnum<MOD> f(int k) {
assert((MOD-1)%k == 0);
return power(modnum<MOD>(primitive_root<MOD>::value), (MOD-1)/k);
}
};
template <> struct getRoot<mod_goldilocks> {
static mod_goldilocks f(int k) {
assert((mod_goldilocks::MOD-1)%k == 0);
return power(mod_goldilocks(mod_goldilocks::PRIMITIVE_ROOT), (mod_goldilocks::MOD-1)/k);
}
};
// We take the bit-reverse convention: the coefficient of a[i] -> b[j] is omega^{i * bit_reverse(j)}.
// This means that the size 2^{k-1} transform is the prefix of the size 2^k transform (wrapping the input).
//
// We mostly work with spans here:
// spans of transforms are expected to have length exactly 2^k
// spans of inputs/outputs are expected to have length [0, 2^{k+1})
//
// Inputs/outputs are treated mod x^{2^k} - 1.
// Their length is allowed to be bigger than 2^k mostly to perform ops on sequences of size n+1 with only transforms of size n.
// The upper bound of 2^{k+1} is arbitrary: we could tighten it to 2^k + 1 or loosen it to infinity, this is just a "defensive" choice.
template <typename num> struct fft_core {
static inline vector<int> rev;
// rt[2^k + i] = 1^{i / 2^(k+1)}
// TODO: can we get rid of inv_rt; alternatively, should we store inv_rt in bit-reverse order?
static inline vector<num> rt, inv_rt;
static void init(int n) {
if (n <= sz(rt)) return;
rev.resize(n);
for (int i = 0; i < n; i++) {
rev[i] = (rev[i>>1] | ((i&1)*n)) >> 1;
}
rt.reserve(n); inv_rt.reserve(n);
while (sz(rt) < 2 && sz(rt) < n) { rt.push_back(num(1)); inv_rt.push_back(num(1)); }
for (int k = sz(rt); k < n; k *= 2) {
rt.resize(2*k); inv_rt.resize(2*k);
num z = getRoot<num>::f(2*k);
num iz = inv(z);
for (int i = k/2; i < k; i++) {
rt[2*i] = rt[i], rt[2*i+1] = rt[i]*z;
inv_rt[2*i] = inv_rt[i], inv_rt[2*i+1] = inv_rt[i]*iz;
}
}
}
// bit-reversal of i as a log2(n)-bit number
static int brev(int i, int n) {
int s = __builtin_ctz(unsigned(sz(rev)/n));
return rev[i] >> s;
}
// index of the conjugate evaluation point, in the returned bit-reversed order
static int conj_index(int j) {
return j == 0 ? 0 : j ^ ((1 << (31 - __builtin_clz(unsigned(j)))) - 1);
}
static void forward(std::span<num> a) {
int n = sz(a);
if (n <= 1) return;
init(n);
for (int k = n/2; k >= 1; k /= 2) {
for (int i = 0; i < n; i += 2*k) {
for (int j = 0; j < k; j++) {
num u = a[i+j], v = a[i+j+k];
a[i+j] = u + v;
a[i+j+k] = (u - v) * rt[j+k];
}
}
}
}
static void inverse(std::span<num> a) {
int n = sz(a);
if (n <= 1) return;
init(n);
for (int k = 1; k < n; k *= 2) {
for (int i = 0; i < n; i += 2*k) {
for (int j = 0; j < k; j++) {
num t = inv_rt[j+k] * a[i+j+k];
a[i+j+k] = a[i+j] - t;
a[i+j] = a[i+j] + t;
}
}
}
}
// Extend a size 2^{k-1} transform to size 2^k; we need the coefficients.
// t must have size 2^k, and coeffs must have size at most 2^{k+1}.
static void extend(std::span<num> t, std::span<const num> coeffs) {
int n = sz(t) / 2;
assert(sz(coeffs) <= 2 * n);
init(sz(t));
auto b = t.subspan(n, n);
int lo = min(sz(coeffs), n);
for (int i = 0; i < lo; i++) {
// rt[n + i] = w_{2n}^i
b[i] = coeffs[i] * rt[n + i];
}
std::fill(b.begin() + lo, b.end(), num(0));
for (int i = n; i < sz(coeffs); i++) {
b[i - n] = b[i - n] - coeffs[i] * rt[i];
}
forward(b);
}
// Consider t = transform(P) and P(x) = E(x^2) + O(x^2) * x
// `even_half` and `odd_half` extract a size 2^{k-1} transform of E/O, respectively.
static void even_half(std::span<const num> t, std::span<num> out) {
int n = sz(out);
assert(sz(t) >= 2*n);
num half = inv(num(2));
for (int j = 0; j < n; j++) out[j] = (t[2*j] + t[2*j+1]) * half;
}
static void odd_half(std::span<const num> t, std::span<num> out) {
int n = sz(out);
assert(sz(t) >= 2*n);
init(2*n);
num half = inv(num(2));
for (int j = 0; j < n; j++) {
// entry j of the size-2n transform pairs (w, -w) with w = w_{2n}^{brev(j, n)}
out[j] = (t[2*j] - t[2*j+1]) * half * inv_rt[n + brev(j, n)];
}
}
};
/* namespace ecnerwala::fft */ }
#line 11 "src/fft/engines/ntt.hpp"
namespace ecnerwala::fft::engines {
template <typename num> struct ntt {
using value_type = num;
static constexpr bool commutative = true;
using core = fft_core<num>;
struct transformed {
vector<num> v;
int size() const { return sz(v); }
};
using product = transformed;
static constexpr int unit_scale = 0;
template <int A = 0> using transformed_t = transformed;
template <int K = 0> using product_t = product;
static transformed transform(std::span<const num> a, int n) {
assert(sz(a) <= 2 * n);
transformed r;
r.v.assign(n, num(0));
int lo = min(sz(a), n);
std::copy(a.begin(), a.begin() + lo, r.v.begin());
for (int i = n; i < sz(a); i++) r.v[i - n] += a[i];
core::forward(std::span<num>(r.v));
return r;
}
static void extend_to(transformed& t, int m, std::span<const num> coeffs) {
assert(!(m & (m-1)) && sz(coeffs) <= 2 * m);
if (t.size() >= m) return;
if (t.size() == 0) { t = transform(coeffs, m); return; }
while (t.size() < m) {
int s = t.size();
t.v.resize(2 * s);
// coeffs past 2s are zero: they didn't fit in the transform we're a prefix of
core::extend(std::span<num>(t.v), coeffs.first(size_t(min(sz(coeffs), 2 * s))));
}
}
static transformed downsample(const transformed& t, int n, bool odd) {
transformed r; r.v.resize(n);
if (odd) core::odd_half(std::span<const num>(t.v), std::span<num>(r.v));
else core::even_half(std::span<const num>(t.v), std::span<num>(r.v));
return r;
}
static transformed negate_arg(const transformed& t, int n) {
assert(n >= 2 && t.size() >= n);
transformed r; r.v.resize(n);
for (int j = 0; j < n; j++) r.v[j] = t.v[j ^ 1];
return r;
}
static product mul(const transformed& a, const transformed& b, int n) {
assert(a.size() >= n && b.size() >= n);
product p; p.v.resize(n);
for (int i = 0; i < n; i++) p.v[i] = a.v[i] * b.v[i];
return p;
}
static product sq(const transformed& a, int n) { return mul(a, a, n); }
static product mul2(
const transformed& a1, const transformed& b1,
const transformed& a2, const transformed& b2,
int n
) {
assert(a1.size() >= n && b1.size() >= n && a2.size() >= n && b2.size() >= n);
product p; p.v.resize(n);
for (int i = 0; i < n; i++) p.v[i] = a1.v[i] * b1.v[i] + a2.v[i] * b2.v[i];
return p;
}
static product add(product&& a, const product& b) {
assert(a.size() == b.size());
for (int i = 0; i < a.size(); i++) a.v[i] += b.v[i];
return std::move(a);
}
template <typename Op = assign_op> static void finish(product&& p, std::span<num> out, Op op = {}) {
int n = p.size();
assert(sz(out) <= n);
core::inverse(std::span<num>(p.v));
num d = inv(num(n));
for (int i = 0; i < sz(out); i++) op(out[i], p.v[i] * d);
}
};
/* namespace ecnerwala::fft::engines */ }
#line 2 "src/fft/engines/real.hpp"
#line 8 "src/fft/engines/real.hpp"
#line 11 "src/fft/engines/real.hpp"
namespace ecnerwala::fft::engines {
// Convolve real (floating point) values by packing into complex numbers with
// a'[t] = a[2t] + i * a[2t+1]
// We use conjugate symmetry to untangle/retangle the two.
// TODO: Add type bounds?
template <typename dbl = double> struct real {
using value_type = dbl;
static constexpr bool commutative = true;
using cnum = cplx<dbl>;
using core = fft_core<cnum>;
struct transformed {
vector<cnum> v;
int size() const { return 2 * sz(v); }
};
using product = transformed;
// Precision is caller-managed for this engine (see add), so scale is untracked.
static constexpr int unit_scale = 0;
template <int A = 0> using transformed_t = transformed;
template <int K = 0> using product_t = product;
static int packed_size(int n) { return std::max(n / 2, 1); }
static void pack(std::span<const dbl> a, std::span<cnum> c) {
for (int i = 0; i < sz(a); i++) (i & 1 ? c[i/2].y : c[i/2].x) = a[i];
}
// Spectrum of the real (odd = false) or imaginary (odd = true) part of the packed
// sequence at bitrev entry t, by conjugate symmetry with the entry of w^{-k}.
static cnum part(const transformed& f, int t, bool odd) {
cnum g = conj(f.v[core::conj_index(t)]);
return odd ? (f.v[t] - g) * cnum(0, dbl(-0.5)) : (f.v[t] + g) * cnum(dbl(0.5));
}
// Given the spectra (s0, s1) of a real sequence x at w_{2mo}^q and w_{2mo}^{q+mo},
// the packed-transform entry of x at packed size mo: the even/odd interleaves of x
// have spectra (s0 +- s1)/2 (the odd one twisted by w_{2mo}^{-q}).
static cnum retangle(cnum s0, cnum s1, int mo, int q) {
cnum s = (s0 + s1) * cnum(dbl(0.5));
cnum d = (s0 - s1) * cnum(dbl(0.5)) * core::inv_rt[mo + q];
return s + cnum(-d.y, d.x);
}
static transformed transform(std::span<const dbl> a, int n) {
assert(sz(a) <= 2 * n);
transformed r;
r.v.assign(packed_size(n), cnum(0));
for (int i = 0; i < sz(a); i++) {
int j = i < n ? i : i - n;
((j & 1) ? r.v[j/2].y : r.v[j/2].x) += a[i];
}
core::forward(std::span<cnum>(r.v));
return r;
}
static void extend_to(transformed& t, int m, std::span<const dbl> coeffs) {
assert(!(m & (m-1)) && sz(coeffs) <= 2 * m);
if (t.size() >= m) return;
if (t.size() == 0) { t = transform(coeffs, m); return; }
auto buf = buffer_pool<cnum>::get((sz(coeffs) + 1) / 2);
std::fill(buf.span().begin(), buf.span().end(), cnum(0));
pack(coeffs, buf.span());
while (t.size() < m) {
int s = sz(t.v);
t.v.resize(2 * s);
// packed coeffs past 2s are zero: they didn't fit in the transform we're a prefix of
core::extend(
std::span<cnum>(t.v),
std::span<const cnum>(buf.span()).first(size_t(min(sz(buf.span()), 2 * s)))
);
}
}
static transformed downsample(const transformed& t, int n, bool odd) { return half(t, n, odd); }
// A(-x) negates the odd (imaginary-slot) coefficients, i.e. conjugates the packed
// sequence; the transform of a conjugated sequence is the conjugate at w^(-k).
static transformed negate_arg(const transformed& t, int n) {
int m = packed_size(n);
assert(n >= 2 && sz(t.v) >= m);
transformed r; r.v.resize(m);
for (int j = 0; j < m; j++) r.v[j] = conj(t.v[core::conj_index(j)]);
return r;
}
static transformed half(const transformed& f, int n, bool odd) {
assert(n >= 2 && f.size() >= 2 * n);
int mo = n / 2;
core::init(2 * mo);
transformed r; r.v.resize(mo);
for (int u = 0; u < mo; u++) {
r.v[u] = retangle(part(f, 2*u, odd), part(f, 2*u+1, odd), mo, core::brev(u, mo));
}
return r;
}
static product mul(const transformed& a, const transformed& b, int n) {
int m = packed_size(n);
assert(a.size() >= n && b.size() >= n);
core::init(2 * m);
product p; p.v.resize(m);
for (int t = 0; t < m; t++) {
int k = core::brev(t, m);
cnum w = core::rt[m + k];
cnum xa = part(a, t, false), ya = part(a, t, true);
cnum xb = part(b, t, false), yb = part(b, t, true);
// full spectra at w_{2m}^k and w_{2m}^{k+m} = -w_{2m}^k
cnum p0 = (xa + w * ya) * (xb + w * yb);
cnum p1 = (xa - w * ya) * (xb - w * yb);
p.v[t] = retangle(p0, p1, m, k);
}
return p;
}
static product sq(const transformed& a, int n) { return mul(a, a, n); }
static product mul2(
const transformed& a1, const transformed& b1,
const transformed& a2, const transformed& b2,
int n
) {
int m = packed_size(n);
assert(a1.size() >= n && b1.size() >= n && a2.size() >= n && b2.size() >= n);
core::init(2 * m);
product p; p.v.resize(m);
for (int t = 0; t < m; t++) {
int k = core::brev(t, m);
cnum w = core::rt[m + k];
cnum xa1 = part(a1, t, false), ya1 = part(a1, t, true);
cnum xb1 = part(b1, t, false), yb1 = part(b1, t, true);
cnum xa2 = part(a2, t, false), ya2 = part(a2, t, true);
cnum xb2 = part(b2, t, false), yb2 = part(b2, t, true);
cnum p0 = (xa1 + w * ya1) * (xb1 + w * yb1) + (xa2 + w * ya2) * (xb2 + w * yb2);
cnum p1 = (xa1 - w * ya1) * (xb1 - w * yb1) + (xa2 - w * ya2) * (xb2 - w * yb2);
p.v[t] = retangle(p0, p1, m, k);
}
return p;
}
static product add(product&& a, const product& b) {
assert(a.size() == b.size());
for (int i = 0; i < sz(a.v); i++) a.v[i] = a.v[i] + b.v[i];
return std::move(a);
}
template <typename Op = assign_op> static void finish(product&& p, std::span<dbl> out, Op op = {}) {
int m = sz(p.v);
assert(sz(out) <= 2 * m);
core::inverse(std::span<cnum>(p.v));
dbl d = dbl(1) / dbl(m);
for (int i = 0; i < sz(out); i++) op(out[i], (i & 1 ? p.v[i/2].y : p.v[i/2].x) * d);
}
};
/* namespace ecnerwala::fft::engines */ }
#line 2 "src/fft/engines/split.hpp"
#line 10 "src/fft/engines/split.hpp"
#line 13 "src/fft/engines/split.hpp"
namespace ecnerwala::fft::engines {
// Multiplies mod `mnum` by splitting values into balanced 15-bit halves (each limb in
// [-2^14, 2^14], from the balanced representative |v| <= MOD/2) packed into one complex
// transform per operand.
template <typename mnum> struct split {
static_assert(sizeof(decltype(mnum::MOD)) <= 4, "limbs must fit 15 bits");
using value_type = mnum;
static constexpr bool commutative = true;
static constexpr int unit_scale = 1;
using cnum = cplx<double>;
using core = fft_core<cnum>;
template <int A = 1> struct transformed_t {
vector<cnum> v;
int size() const { return sz(v); }
transformed_t() = default;
explicit transformed_t(vector<cnum>&& v_) : v(std::move(v_)) {}
template <int A2> requires (A2 != A) explicit(A2 > A) transformed_t(transformed_t<A2>&& o)
: v(std::move(o.v)) {}
};
using transformed = transformed_t<1>;
template <int K> struct product_t {
// After finish's inverse transforms: lo = (lo*lo, hi*lo), hi = (lo*hi, hi*hi).
vector<cnum> lo, hi;
int size() const { return sz(lo); }
product_t() = default;
product_t(vector<cnum>&& lo_, vector<cnum>&& hi_) : lo(std::move(lo_)), hi(std::move(hi_)) {}
template <int K2> requires (K2 != K) explicit(K2 > K) product_t(product_t<K2>&& o)
: lo(std::move(o.lo)), hi(std::move(o.hi)) {}
};
using product = product_t<1>;
static cnum pack(mnum x) {
int64_t v = x.balanced();
int64_t hi = (v + (1 << 14)) >> 15;
return cnum(double(v - (hi << 15)), double(hi));
}
static transformed transform(std::span<const mnum> a, int n) {
assert(sz(a) <= 2 * n);
transformed r;
r.v.assign(n, cnum(0));
for (int i = 0; i < sz(a); i++) {
int j = i < n ? i : i - n;
r.v[j] = r.v[j] + pack(a[i]);
}
core::forward(std::span<cnum>(r.v));
return r;
}
static void extend_to(transformed& t, int m, std::span<const mnum> coeffs) {
assert(!(m & (m-1)) && sz(coeffs) <= 2 * m);
if (t.size() >= m) return;
if (t.size() == 0) { t = transform(coeffs, m); return; }
auto buf = buffer_pool<cnum>::get(sz(coeffs));
for (int i = 0; i < sz(coeffs); i++) buf[i] = pack(coeffs[i]);
while (t.size() < m) {
int s = t.size();
t.v.resize(2 * s);
// coeffs past 2s are zero: they didn't fit in the transform we're a prefix of
core::extend(
std::span<cnum>(t.v),
std::span<const cnum>(buf.span()).first(size_t(min(sz(coeffs), 2 * s)))
);
}
}
static void downsample_core(std::span<const cnum> in, std::span<cnum> out, bool odd) {
if (odd) core::odd_half(in, out);
else core::even_half(in, out);
}
template <int A> static transformed_t<A> downsample(const transformed_t<A>& t, int n, bool odd) {
transformed_t<A> r; r.v.resize(n);
downsample_core(std::span<const cnum>(t.v), std::span<cnum>(r.v), odd);
return r;
}
template <int K> static product_t<K> downsample(const product_t<K>& p, int n, bool odd) {
product_t<K> r; r.lo.resize(n); r.hi.resize(n);
downsample_core(std::span<const cnum>(p.lo), std::span<cnum>(r.lo), odd);
downsample_core(std::span<const cnum>(p.hi), std::span<cnum>(r.hi), odd);
return r;
}
template <int A> static transformed_t<A> negate_arg(const transformed_t<A>& t, int n) {
assert(n >= 2 && t.size() >= n);
transformed_t<A> r; r.v.resize(n);
for (int j = 0; j < n; j++) r.v[j] = t.v[j ^ 1];
return r;
}
template <int A, int B> static transformed_t<A + B> add(transformed_t<A>&& a, const transformed_t<B>& b) {
transformed_t<A + B> r{std::move(a.v)};
add_into(r.v, b.v);
return r;
}
// Unpacks b's transform into transforms of its low/high halves via conjugate
// symmetry, then multiplies both against a's (still packed) transform. The scale
// parameter only affects the bookkeeping, so the body is a shared untyped impl.
static void mul_impl(const vector<cnum>& a, const vector<cnum>& b, vector<cnum>& lo, vector<cnum>& hi, int n, bool acc = false) {
core::init(n);
lo.resize(n); hi.resize(n);
for (int i = 0; i < n; i++) {
int ci = core::conj_index(i);
cnum g0 = (b[i] + conj(b[ci])) * cnum(0.5);
cnum t = (b[i] - conj(b[ci])) * cnum(0.5);
cnum g1 = cnum(t.y, -t.x);
if (acc) {
lo[i] = lo[i] + a[i] * g0;
hi[i] = hi[i] + a[i] * g1;
} else {
lo[i] = a[i] * g0;
hi[i] = a[i] * g1;
}
}
}
template <int A, int B> static product_t<A * B> mul(const transformed_t<A>& a, const transformed_t<B>& b, int n) {
assert(a.size() >= n && b.size() >= n);
product_t<A * B> p;
mul_impl(a.v, b.v, p.lo, p.hi, n);
return p;
}
template <int A> static product_t<A * A> sq(const transformed_t<A>& a, int n) { return mul(a, a, n); }
template <int A1, int B1, int A2, int B2>
static product_t<A1 * B1 + A2 * B2> mul2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int n
) {
assert(a1.size() >= n && b1.size() >= n && a2.size() >= n && b2.size() >= n);
product_t<A1 * B1 + A2 * B2> p;
mul_impl(a1.v, b1.v, p.lo, p.hi, n);
mul_impl(a2.v, b2.v, p.lo, p.hi, n, true);
return p;
}
static void add_into(vector<cnum>& a, const vector<cnum>& b) {
assert(sz(a) == sz(b));
for (int i = 0; i < sz(a); i++) a[i] = a[i] + b[i];
}
template <int K1, int K2> static product_t<K1 + K2> add(product_t<K1>&& a, product_t<K2>&& b) {
product_t<K1 + K2> r{std::move(a.lo), std::move(a.hi)};
add_into(r.lo, b.lo);
add_into(r.hi, b.hi);
return r;
}
template <int K = 1, typename Op = assign_op> static void finish(product_t<K>&& p, std::span<mnum> out, Op op = {}) {
// The fp error budget is divided by the accumulated scale; K <= 2 is very
// conservative (balanced limbs already left ~2x headroom at max lengths).
static_assert(K <= 2, "split: accumulated scale too large");
int n = p.size();
assert(sz(out) <= n);
core::inverse(std::span<cnum>(p.lo));
core::inverse(std::span<cnum>(p.hi));
const int64_t m = mnum::MOD;
double d = 1.0 / double(n);
// llround + a final wrap so negative half-products (e.g. from negate_arg'd
// transforms) reconstruct correctly.
for (int i = 0; i < sz(out); i++) {
int64_t v = (llround(p.lo[i].x * d)
+ (llround(p.lo[i].y * d) % m << 15)
+ (llround(p.hi[i].x * d) % m << 15)
+ (llround(p.hi[i].y * d) % m << 30)) % m;
if (v < 0) v += m;
op(out[i], mnum(v));
}
}
};
/* namespace ecnerwala::fft::engines */ }
#line 2 "src/fft/engines/crt.hpp"
#line 7 "src/fft/engines/crt.hpp"
#line 11 "src/fft/engines/crt.hpp"
namespace ecnerwala::fft::engines {
// Multiplies mod `mnum` by running NTTs modulo two FFT-friendly primes and CRT'ing.
// Inputs use balanced representatives (|v| <= MOD/2), so the true integer coefficients
// are bounded by n (MOD/2)^2.
template <typename mnum, typename num1 = mod_goldilocks, typename num2 = modnum<(15 << 27) + 1>>
struct crt {
static_assert(sizeof(decltype(mnum::MOD)) <= 4, "n (MOD/2)^2 must fit the CRT modulus product");
using value_type = mnum;
static constexpr bool commutative = true;
static constexpr int unit_scale = 1;
using E1 = ntt<num1>;
using E2 = ntt<num2>;
template <int A = 1> struct transformed_t {
typename E1::transformed t1;
typename E2::transformed t2;
int size() const { return t1.size(); }
transformed_t() = default;
transformed_t(typename E1::transformed&& t1_, typename E2::transformed&& t2_)
: t1(std::move(t1_)), t2(std::move(t2_)) {}
template <int A2> requires (A2 != A) explicit(A2 > A) transformed_t(transformed_t<A2>&& o)
: t1(std::move(o.t1)), t2(std::move(o.t2)) {}
};
using transformed = transformed_t<1>;
template <int K> struct product_t {
typename E1::product p1;
typename E2::product p2;
int size() const { return sz(p1); }
product_t() = default;
product_t(typename E1::product&& p1_, typename E2::product&& p2_)
: p1(std::move(p1_)), p2(std::move(p2_)) {}
template <int K2> requires (K2 != K) explicit(K2 > K) product_t(product_t<K2>&& o)
: p1(std::move(o.p1)), p2(std::move(o.p2)) {}
};
using product = product_t<1>;
static transformed transform(std::span<const mnum> a, int n) {
assert(sz(a) <= 2 * n);
auto b1 = buffer_pool<num1>::get(sz(a));
auto b2 = buffer_pool<num2>::get(sz(a));
for (int i = 0; i < sz(a); i++) { int64_t v = a[i].balanced(); b1[i] = num1(v); b2[i] = num2(v); }
return transformed{
E1::transform(std::span<const num1>(b1.span()), n),
E2::transform(std::span<const num2>(b2.span()), n),
};
}
static void extend_to(transformed& t, int m, std::span<const mnum> coeffs) {
if (t.size() >= m) return;
auto b1 = buffer_pool<num1>::get(sz(coeffs));
auto b2 = buffer_pool<num2>::get(sz(coeffs));
for (int i = 0; i < sz(coeffs); i++) { int64_t v = coeffs[i].balanced(); b1[i] = num1(v); b2[i] = num2(v); }
E1::extend_to(t.t1, m, std::span<const num1>(b1.span()));
E2::extend_to(t.t2, m, std::span<const num2>(b2.span()));
}
template <int A> static transformed_t<A> downsample(const transformed_t<A>& t, int n, bool odd) {
return transformed_t<A>{E1::downsample(t.t1, n, odd), E2::downsample(t.t2, n, odd)};
}
template <int K> static product_t<K> downsample(const product_t<K>& p, int n, bool odd) {
return product_t<K>{E1::downsample(p.p1, n, odd), E2::downsample(p.p2, n, odd)};
}
template <int A> static transformed_t<A> negate_arg(const transformed_t<A>& t, int n) {
return transformed_t<A>{E1::negate_arg(t.t1, n), E2::negate_arg(t.t2, n)};
}
// Exact per prime; the scale tracks the true (integer) coefficient growth.
template <int A, int B> static transformed_t<A + B> add(transformed_t<A>&& a, const transformed_t<B>& b) {
return transformed_t<A + B>{E1::add(std::move(a.t1), b.t1), E2::add(std::move(a.t2), b.t2)};
}
template <int A, int B> static product_t<A * B> mul(const transformed_t<A>& a, const transformed_t<B>& b, int n) {
return product_t<A * B>{E1::mul(a.t1, b.t1, n), E2::mul(a.t2, b.t2, n)};
}
template <int A> static product_t<A * A> sq(const transformed_t<A>& a, int n) { return mul(a, a, n); }
template <int A1, int B1, int A2, int B2>
static product_t<A1 * B1 + A2 * B2> mul2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int n
) {
return product_t<A1 * B1 + A2 * B2>{
E1::mul2(a1.t1, b1.t1, a2.t1, b2.t1, n),
E2::mul2(a1.t2, b1.t2, a2.t2, b2.t2, n),
};
}
template <int K1, int K2> static product_t<K1 + K2> add(product_t<K1>&& a, product_t<K2>&& b) {
return product_t<K1 + K2>{E1::add(std::move(a.p1), b.p1), E2::add(std::move(a.p2), b.p2)};
}
template <int K = 1, typename Op = assign_op> static void finish(product_t<K>&& p, std::span<mnum> out, Op op = {}) {
// The reconstruction needs |c| < whole/2; balanced inputs bound each addend's
// true coefficients by n (MOD/2)^2, so the safe length is divided by the
// accumulated scale. K <= 2 is very conservative (~2^35 even for MOD ~ 2^30).
static_assert(K <= 2, "crt: accumulated scale too large");
int n = p.size();
assert(sz(out) <= n);
auto o1 = buffer_pool<num1>::get(sz(out));
auto o2 = buffer_pool<num2>::get(sz(out));
E1::finish(std::move(p.p1), o1.span());
E2::finish(std::move(p.p2), o2.span());
// TODO: Could hardcode these
num1 inv_n2 = inv(num1(num2::MOD));
num2 inv_n1 = inv(num2(num1::MOD));
__int128_t whole = __int128_t(num1::MOD) * __int128_t(num2::MOD);
mnum m1_mod = mnum(num1::MOD);
mnum m2_mod = mnum(num2::MOD);
mnum whole_mod = m1_mod * m2_mod;
for (int i = 0; i < sz(out); i++) {
num1 v1 = o1[i] * inv_n2;
num2 v2 = o2[i] * inv_n1;
mnum o_mod = mnum(uint64_t(v1)) * m2_mod + mnum(int(v2)) * m1_mod;
__int128_t o_exact = __int128_t(uint64_t(v1)) * __int128_t(num2::MOD) + __int128_t(int(v2)) * __int128_t(num1::MOD);
if (o_exact >= whole) { o_exact -= whole; o_mod -= whole_mod; }
// Balanced representatives: |o| <= whole/2
if (o_exact > whole / 2) o_mod -= whole_mod;
op(out[i], o_mod);
}
}
};
/* namespace ecnerwala::fft::engines */ }
#line 11 "src/fft/test_util.test.hpp"
// Shared helpers for the fft/ unit tests.
namespace ecnerwala {
namespace fft {
template <typename T> std::vector<T> multiply_slow(const std::vector<T>& a, const std::vector<T>& b) {
if (a.empty() || b.empty()) return {};
std::vector<T> res(a.size() + b.size() - 1);
for (int i = 0; i < int(a.size()); i++) {
for (int j = 0; j < int(b.size()); j++) {
res[i+j] += a[i] * b[j];
}
}
return res;
}
// Small values for doubles so products are exactly representable; full range otherwise.
template <typename T> T rnd_val(std::mt19937& mt) {
if constexpr (std::is_floating_point_v<T>) return T(int(mt() % 1024));
else return T(mt());
}
template <typename T> void fill_rnd(std::vector<T>& v, std::mt19937& mt) {
for (T& x : v) x = rnd_val<T>(mt);
}
template <typename T> void check_eq(const std::vector<T>& got, const std::vector<T>& want) {
REQUIRE(got.size() == want.size());
for (int i = 0; i < int(got.size()); i++) {
INFO("i = " << i);
if constexpr (std::is_floating_point_v<T>) REQUIRE(llround(got[i]) == llround(want[i]));
else REQUIRE(got[i] == want[i]);
}
}
#define ALL_ENGINES \
engines::ntt<modnum<998244353>>, engines::ntt<mod_goldilocks>, engines::real<double>, \
engines::split<modnum<int(1e9)+7>>, engines::crt<modnum<int(1e9)+7>>
#define MOD_ENGINES \
engines::ntt<modnum<998244353>>, engines::ntt<mod_goldilocks>, \
engines::split<modnum<int(1e9)+7>>, engines::crt<modnum<int(1e9)+7>>
}} // namespace ecnerwala::fft
#line 9 "src/fft/poly.test.cpp"
namespace ecnerwala {
namespace fft {
using namespace std;
TEST_CASE("poly::form evaluation and transposed multiplication", "[fft]") {
using num = modnum<998244353>;
using E = engines::ntt<num>;
mt19937 mt(Catch::getSeed());
int n = 40;
num z = num(mt());
auto f = poly::form<E>::polynomial_evaluation(z, n);
vector<num> sv(30);
fill_rnd(sv, mt);
poly::vec<E> s((span<const num>(sv)));
REQUIRE(f(s) == s(z));
// f(S * q) == composed_with(q)(S)
vector<num> qv(11), s2v(n - 11);
fill_rnd(qv, mt);
fill_rnd(s2v, mt);
poly::vec<E> q((span<const num>(qv))), s2((span<const num>(s2v)));
auto fq = f.composed_with(q);
REQUIRE(fq(s2) == f(s2 * q));
// evaluation functional composed with q evaluates S * q at z
REQUIRE(fq(s2) == (s2 * q)(z));
// composing with a power series (living in 1/x) multiplies the storages,
// prefix-truncated back to length n
series::trunc<E> t(size_t(n), num{});
fill_rnd(t, mt);
auto ft = f.composed_with(t);
REQUIRE(ft.len() == n);
for (int j = 0; j < n; j++) {
num want{};
for (int d = 0; d <= j; d++) want += t[d] * f.rev_series()[j - d];
REQUIRE(ft.rev_series()[j] == want);
}
// an exact series may be any length; the tail beyond it is zero
series::exact<E> e(t.begin(), t.begin() + 11);
auto fe = f.composed_with(e);
REQUIRE(fe.len() == n);
for (int j = 0; j < n; j++) {
num want{};
for (int d = 0; d <= j && d < 11; d++) want += t[d] * f.rev_series()[j - d];
REQUIRE(fe.rev_series()[j] == want);
}
}
TEST_CASE("poly::vec reversed storage and series interop", "[fft]") {
using num = modnum<998244353>;
using E = engines::ntt<num>;
mt19937 mt(Catch::getSeed());
vector<num> pa(37), pb(23);
fill_rnd(pa, mt);
fill_rnd(pb, mt);
poly::vec<E> a((span<const num>(pa))), b((span<const num>(pb)));
// indexing is coefficient order, storage is reversed
REQUIRE(a[0] == pa[0]);
REQUIRE(a.leading() == pa[36]);
REQUIRE(a.rev_series()[0] == pa[36]);
REQUIRE(a.rev_series()[36] == pa[0]);
// products convolve the storage directly
poly::vec<E> p = a * b;
check_eq(vector<num>(p.begin(), p.end()), multiply_slow(pa, pb));
REQUIRE(square(a) == a * a);
num x = num(mt());
REQUIRE(p(x) == a(x) * b(x));
// +/- align at x^0 (the shared storage tail)
poly::vec<E> s = a + b, d = b - a;
REQUIRE(s.len() == 37);
for (int i = 0; i < 37; i++) REQUIRE(s[i] == pa[i] + (i < 23 ? pb[i] : num(0)));
for (int i = 0; i < 37; i++) REQUIRE(d[i] == (i < 23 ? pb[i] : num(0)) - pa[i]);
// multiplying by x^k appends zeros to the storage; coefficients shift up
poly::vec<E> g = a;
g.shift(2);
REQUIRE(g.len() == 39);
REQUIRE(g[0] == num(0));
REQUIRE(g[1] == num(0));
for (int i = 0; i < 37; i++) REQUIRE(g[i + 2] == pa[i]);
REQUIRE(g.rev_series().data()[0] == pa[36]);
// named conversions use the reversed convention and round-trip freely
const series::exact<E>& ra = a.rev_series();
REQUIRE(ra.len() == 37);
for (int i = 0; i < 37; i++) REQUIRE(ra[i] == pa[36 - i]);
REQUIRE(poly::vec<E>::from_rev_series(ra) == a);
// a poly::vec's natural-order coefficients as an exact series (for series products)
series::exact<E> xa(a.begin(), a.end());
REQUIRE(equal(xa.begin(), xa.end(), pa.begin(), pa.end()));
REQUIRE(a.unrev_series(10) == series::trunc<E>(pa.begin(), pa.begin() + 10));
// the storage transform serves transposed products: middle product against rev_series()
std::vector<num> vals(60);
fill_rnd(vals, mt);
series::cached_exact<E> cv(series::exact<E>(vals.begin(), vals.end()));
series::cached_exact<E> ca(a.rev_series());
auto mp = middle_product(cv, ca);
auto naive = [&](int j) {
num r{};
for (int t = 0; t < 37; t++) r += pa[t] * vals[j + t];
return r;
};
for (int j = 0; j < sz(mp); j++) REQUIRE(mp[size_t(j)] == naive(j));
}
TEST_CASE("poly::cached products", "[fft]") {
using num = modnum<998244353>;
using E = engines::ntt<num>;
mt19937 mt(Catch::getSeed());
vector<num> pa(37), pb(23);
fill_rnd(pa, mt);
fill_rnd(pb, mt);
poly::vec<E> a((span<const num>(pa))), b((span<const num>(pb)));
// poly::vec products return poly::cached, adopting the product transform
auto p = a * b;
static_assert(std::is_same_v<decltype(p), poly::cached<E>>);
REQUIRE(p.rev_series().cache().size() > 0);
poly::vec<E> pp = a * b; // naming the plain type moves out and drops the transform
check_eq(vector<num>(pp.begin(), pp.end()), multiply_slow(pa, pb));
// cached operands reuse and chain; results compare across representations
poly::cached<E> ca(a), cb(b);
REQUIRE(ca == a);
REQUIRE(ca * cb == p);
REQUIRE(ca * b == p);
auto sq = square(ca);
static_assert(std::is_same_v<decltype(sq), poly::cached<E>>);
REQUIRE(sq == a * a);
num x = num(mt());
REQUIRE(p(x) == ca(x) * cb(x));
// moving out drops down to a plain mutable poly::vec
poly::vec<E> q = std::move(ca);
REQUIRE(q == a);
}
TEST_CASE("poly::multipoint and poly::interpolate", "[fft]") {
using num = modnum<998244353>;
mt19937 mt(Catch::getSeed());
for (int n : {1, 2, 3, 8, 17, 40}) {
INFO("n = " << n);
vector<num> coeffs(n);
for (num& x : coeffs) { x = num(mt()); }
poly::vec<engines::ntt<num>> p((span<const num>(coeffs)));
vector<num> pts(n);
for (int i = 0; i < n; i++) pts[i] = num(1000 + i);
auto vals = poly::multipoint<engines::ntt<num>>(p, pts);
for (int i = 0; i < n; i++) {
REQUIRE(vals[i] == p(pts[i]));
}
auto rec = poly::interpolate<engines::ntt<num>>(pts, vals);
REQUIRE(rec == p);
}
}
}} // namespace ecnerwala::fft
// clang-format off
// @formatter:off
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wpragmas"
#pragma GCC diagnostic ignored "-Wunknown-warning-option"
#pragma GCC diagnostic ignored "-Wmisleading-indentation"
#pragma GCC diagnostic ignored "-Wmultistatement-macros"
#include <bits/stdc++.h>
#include <catch2/catch_test_macros.hpp>
#include <catch2/catch_get_random_seed.hpp>
// src/fft/common.hpp
namespace ecnerwala{
template<class T>int sz(T&&arg){using std::size;return int(size(std::forward<T>(arg)));}
inline int nextPow2(int s){return 1<<(s>1?32-__builtin_clz(s-1):0);}
namespace fft{
using std::swap;
using std::vector;
using std::min;
using std::max;
#ifndef ECNERWALA_FFT_POOL_STORAGE
#define ECNERWALA_FFT_POOL_STORAGE
#endif
template<typename T>struct buffer_pool{
static inline ECNERWALA_FFT_POOL_STORAGE std::vector<std::vector<T>>free_list;
struct handle{
std::vector<T>v;
explicit handle(int n){
if(!free_list.empty()){
v=std::move(free_list.back());
free_list.pop_back();
}
v.assign(n,T());
}
handle(const handle&)=delete;
handle&operator=(const handle&)=delete;
handle(handle&&o)noexcept:v(std::move(o.v)){}
~handle(){if(v.capacity())free_list.push_back(std::move(v));}
T&operator[](int i){return v[i];}
operator std::span<T>(){return std::span<T>(v);}
std::span<T>span(){return std::span<T>(v);}
};
static handle get(int n){return handle(n);}
};
}
}
// src/fft/engine.hpp
namespace ecnerwala::fft{
struct assign_op{template<typename T>void operator()(T&d,T v)const{d=v;}};
struct add_op{template<typename T>void operator()(T&d,T v)const{d+=v;}};
struct sub_op{template<typename T>void operator()(T&d,T v)const{d-=v;}};
struct add_twice_op{template<typename T>void operator()(T&d,T v)const{d+=v+v;}};
template<typename E>
concept engine=requires(
std::span<const typename E::value_type>in,
std::span<typename E::value_type>out,
typename E::transformed&t,
const typename E::transformed&ct,
typename E::product&p,
const typename E::product&cp,
int n
){
typename E::value_type;
{E::transform(in,n)}->std::same_as<typename E::transformed>;
{ct.size()}->std::same_as<int>;
E::extend_to(t,n,in);
{E::downsample(ct,n,false)}->std::same_as<typename E::transformed>;
{E::downsample(cp,n,false)}->std::same_as<typename E::product>;
{E::negate_arg(ct,n)}->std::same_as<typename E::transformed>;
{E::mul(ct,ct,n)}->std::same_as<typename E::product>;
{E::sq(ct,n)}->std::same_as<typename E::product>;
{E::mul2(ct,ct,ct,ct,n)}->std::same_as<typename E::template product_t<2*E::unit_scale>>;
E::finish(std::move(p),out);
E::finish(std::move(p),out,add_op{});
E::finish(E::add(std::move(p),std::move(p)),out);
{E::add(E::transform(in,n),ct)}->std::same_as<typename E::template transformed_t<2*E::unit_scale>>;
{E::add(std::move(p),std::move(p))}->std::same_as<typename E::template product_t<2*E::unit_scale>>;
requires std::same_as<std::remove_cvref_t<decltype(E::commutative)>,bool>;
requires std::same_as<std::remove_cvref_t<decltype(E::unit_scale)>,int>;
};
template<typename A,typename B>
concept same_engine=std::same_as<typename A::engine_t,typename B::engine_t>;
template<engine E>using transformed=typename E::transformed;
}
// src/fft/multiply.hpp
namespace ecnerwala::fft{
template<engine E,typename Op=assign_op>
void multiply_circular(std::span<const typename E::value_type>a,std::span<const typename E::value_type>b,
std::span<typename E::value_type>out,int n,Op op={}){
assert(!(n&(n-1)));
auto ta=E::transform(a,n);
auto tb=E::transform(b,n);
E::finish(E::mul(ta,tb,n),out,op);
}
template<engine E,typename Op=assign_op>
void square_circular(std::span<const typename E::value_type>a,std::span<typename E::value_type>out,int n,Op op={}){
assert(!(n&(n-1)));
auto ta=E::transform(a,n);
E::finish(E::sq(ta,n),out,op);
}
namespace detail{
struct conv_size{int n;bool cut;};
inline conv_size conv_size_for(int s){
int n=nextPow2(s);
bool cut=(n==2*(s-1));
return{cut?n/2:n,cut};
}
template<typename T,typename Op>
void emit_linear(std::span<T>buf,int n,int s,bool cut,T c0,std::span<T>out,Op op){
T cn{};
if(cut){
cn=buf[0]-c0;
buf[0]=c0;
}
int lim=min(sz(out),min(s,n));
for(int i=0;i<lim;i++)op(out[i],buf[i]);
if(cut&&sz(out)>=s)op(out[s-1],cn);
}
template<typename T,typename Op>
struct cut_op{
Op op;
T*out0;
T c0;
T&cn;
void operator()(T&x,T v)const{
if(&x==out0){cn=v-c0;v=c0;}
op(x,v);
}
};
template<engine E,typename P,typename Op=assign_op>
void finish_linear(
P&&p,int n,int s,bool cut,
typename E::value_type c0,std::span<typename E::value_type>out,Op op={}
){
using T=typename E::value_type;
if(sz(out)==0)return;
int lim=min(sz(out),min(s,n));
if(!cut){
E::finish(std::move(p),out.subspan(0,lim),op);
}else{
T cn{};
E::finish(std::move(p),out.subspan(0,lim),cut_op<T,Op>{op,&out[0],c0,cn});
if(sz(out)>=s)op(out[s-1],cn);
}
}
}
template<engine E,typename Op=assign_op>
void multiply(std::span<const typename E::value_type>a,std::span<const typename E::value_type>b,
std::span<typename E::value_type>out,Op op={}){
using T=typename E::value_type;
if(sz(a)==0||sz(b)==0)return;
int s=sz(a)+sz(b)-1;
auto[n,cut]=detail::conv_size_for(s);
T c0=a[0]*b[0];
auto buf=buffer_pool<T>::get(n);
multiply_circular<E>(a,b,buf.span(),n);
detail::emit_linear<T>(buf.span(),n,s,cut,c0,out,op);
}
template<engine E,typename Op=assign_op>
void multiply(std::span<const typename E::value_type>a,transformed<E>&ta,
std::span<const typename E::value_type>b,transformed<E>&tb,
std::span<typename E::value_type>out,Op op={}){
using T=typename E::value_type;
if(sz(a)==0||sz(b)==0)return;
int s=sz(a)+sz(b)-1;
auto[n,cut]=detail::conv_size_for(s);
T c0=a[0]*b[0];
E::extend_to(ta,n,a);
E::extend_to(tb,n,b);
detail::finish_linear<E>(E::mul(ta,tb,n),n,s,cut,c0,out,op);
}
template<engine E,typename Op=assign_op>
void multiply_add2(std::span<const typename E::value_type>a1,transformed<E>&ta1,
std::span<const typename E::value_type>b1,transformed<E>&tb1,
std::span<const typename E::value_type>a2,transformed<E>&ta2,
std::span<const typename E::value_type>b2,transformed<E>&tb2,
std::span<typename E::value_type>out,Op op={}){
using T=typename E::value_type;
assert(sz(a1)>0&&sz(b1)>0&&sz(a2)>0&&sz(b2)>0);
int s=sz(a1)+sz(b1)-1;
assert(sz(a2)+sz(b2)-1==s);
auto[n,cut]=detail::conv_size_for(s);
T c0=a1[0]*b1[0]+a2[0]*b2[0];
E::extend_to(ta1,n,a1);E::extend_to(tb1,n,b1);
E::extend_to(ta2,n,a2);E::extend_to(tb2,n,b2);
detail::finish_linear<E>(E::mul2(ta1,tb1,ta2,tb2,n),n,s,cut,c0,out,op);
}
template<engine E>
void multiply_add2_cached(
std::span<const typename E::value_type>a1,transformed<E>&ta1,
std::span<const typename E::value_type>b1,transformed<E>&tb1,
std::span<const typename E::value_type>a2,transformed<E>&ta2,
std::span<const typename E::value_type>b2,transformed<E>&tb2,
std::vector<typename E::value_type>&coeffs,transformed<E>&t){
using T=typename E::value_type;
assert(sz(a1)>0&&sz(b1)>0&&sz(a2)>0&&sz(b2)>0);
int s=sz(a1)+sz(b1)-1;
assert(sz(a2)+sz(b2)-1==s);
coeffs.assign(size_t(s),T{});
t=transformed<E>{};
if constexpr(std::same_as<typename E::product,transformed<E>>){
auto[n,cut]=detail::conv_size_for(s);
T c0=a1[0]*b1[0]+a2[0]*b2[0];
E::extend_to(ta1,n,a1);E::extend_to(tb1,n,b1);
E::extend_to(ta2,n,a2);E::extend_to(tb2,n,b2);
auto p=E::mul2(ta1,tb1,ta2,tb2,n);
auto tp=p;
detail::finish_linear<E>(std::move(p),n,s,cut,c0,std::span<T>(coeffs));
t=std::move(tp);
}else{
multiply_add2<E>(a1,ta1,b1,tb1,a2,ta2,b2,tb2,std::span<T>(coeffs));
}
}
template<engine E>
void multiply_cached(std::span<const typename E::value_type>a,transformed<E>&ta,
std::span<const typename E::value_type>b,transformed<E>&tb,
std::vector<typename E::value_type>&coeffs,transformed<E>&t){
using T=typename E::value_type;
coeffs.assign(size_t(sz(a)&&sz(b)?sz(a)+sz(b)-1:0),T{});
t=transformed<E>{};
if(coeffs.empty())return;
int s=sz(coeffs);
if constexpr(std::same_as<typename E::product,transformed<E>>){
auto[n,cut]=detail::conv_size_for(s);
T c0=a[0]*b[0];
E::extend_to(ta,n,a);
E::extend_to(tb,n,b);
auto p=E::mul(ta,tb,n);
auto tp=p;
detail::finish_linear<E>(std::move(p),n,s,cut,c0,std::span<T>(coeffs));
t=std::move(tp);
}else{
multiply<E>(a,ta,b,tb,std::span<T>(coeffs));
}
}
template<engine E,typename Op=assign_op>
void square(std::span<const typename E::value_type>a,std::span<typename E::value_type>out,Op op={}){
using T=typename E::value_type;
if(sz(a)==0)return;
int s=2*sz(a)-1;
auto[n,cut]=detail::conv_size_for(s);
T c0=a[0]*a[0];
auto buf=buffer_pool<T>::get(n);
square_circular<E>(a,buf.span(),n);
detail::emit_linear<T>(buf.span(),n,s,cut,c0,out,op);
}
template<engine E,typename Op=assign_op>
void square(std::span<const typename E::value_type>a,transformed<E>&ta,
std::span<typename E::value_type>out,Op op={}){
using T=typename E::value_type;
if(sz(a)==0)return;
int s=2*sz(a)-1;
auto[n,cut]=detail::conv_size_for(s);
T c0=a[0]*a[0];
E::extend_to(ta,n,a);
detail::finish_linear<E>(E::sq(ta,n),n,s,cut,c0,out,op);
}
template<engine E>
void square_cached(std::span<const typename E::value_type>a,transformed<E>&ta,
std::vector<typename E::value_type>&coeffs,transformed<E>&t){
using T=typename E::value_type;
coeffs.assign(size_t(sz(a)?2*sz(a)-1:0),T{});
t=transformed<E>{};
if(coeffs.empty())return;
int s=sz(coeffs);
if constexpr(std::same_as<typename E::product,transformed<E>>){
auto[n,cut]=detail::conv_size_for(s);
T c0=a[0]*a[0];
E::extend_to(ta,n,a);
auto p=E::sq(ta,n);
auto tp=p;
detail::finish_linear<E>(std::move(p),n,s,cut,c0,std::span<T>(coeffs));
t=std::move(tp);
}else{
square<E>(a,ta,std::span<T>(coeffs));
}
}
template<engine E>vector<typename E::value_type>multiply(
const vector<typename E::value_type>&a,const vector<typename E::value_type>&b){
using T=typename E::value_type;
if(sz(a)==0||sz(b)==0)return{};
vector<T>r(sz(a)+sz(b)-1);
multiply<E>(std::span<const T>(a),std::span<const T>(b),std::span<T>(r));
return r;
}
template<engine E>vector<typename E::value_type>square(const vector<typename E::value_type>&a){
using T=typename E::value_type;
if(sz(a)==0)return{};
vector<T>r(2*sz(a)-1);
square<E>(std::span<const T>(a),std::span<T>(r));
return r;
}
namespace detail{
template<typename T,typename Op>
void emit_middle(std::span<T>buf,bool cut,int la,int lb,T c0,T ctop,std::span<T>out,Op op){
int m=la-lb+1;
T cn{};
if(cut){
cn=buf[0]-c0;
buf[lb-1]-=ctop;
}
int lim=min(sz(out),cut?m-1:m);
for(int t=0;t<lim;t++)op(out[t],buf[lb-1+t]);
if(cut&&sz(out)>=m)op(out[m-1],cn);
}
}
template<engine E,typename Op=assign_op>
void middle_product(std::span<const typename E::value_type>a,std::span<const typename E::value_type>b,
std::span<typename E::value_type>out,Op op={}){
using T=typename E::value_type;
if(sz(a)==0||sz(b)==0)return;
assert(sz(a)>=sz(b));
if(sz(a)==sz(b)){
T r{};
for(int i=0;i<sz(a);i++){
r+=a[i]*b[sz(b)-1-i];
}
if(sz(out)>0)op(out[0],r);
return;
}
auto[n,cut]=detail::conv_size_for(sz(a));
auto buf=buffer_pool<T>::get(n);
multiply_circular<E>(a,b,buf.span(),n);
detail::emit_middle<T>(buf.span(),cut,sz(a),sz(b),
a[0]*b[0],a[sz(a)-1]*b[sz(b)-1],out,op);
}
template<engine E>vector<typename E::value_type>middle_product(
std::span<const typename E::value_type>a,std::span<const typename E::value_type>b){
using T=typename E::value_type;
if(sz(a)==0||sz(b)==0)return{};
assert(sz(a)>=sz(b));
vector<T>r(size_t(sz(a)-sz(b)+1));
middle_product<E>(a,b,std::span<T>(r));
return r;
}
template<engine E,typename Op=assign_op>
void middle_product(std::span<const typename E::value_type>a,transformed<E>&ta,
std::span<const typename E::value_type>b,transformed<E>&tb,
std::span<typename E::value_type>out,Op op={}){
using T=typename E::value_type;
if(sz(a)==0||sz(b)==0)return;
assert(sz(a)>=sz(b));
if(sz(a)==sz(b)){
T r{};
for(int i=0;i<sz(a);i++){
r+=a[i]*b[sz(b)-1-i];
}
if(sz(out)>0)op(out[0],r);
return;
}
auto[n,cut]=detail::conv_size_for(sz(a));
E::extend_to(ta,n,a);
E::extend_to(tb,n,b);
auto buf=buffer_pool<T>::get(n);
E::finish(E::mul(ta,tb,n),buf.span());
detail::emit_middle<T>(buf.span(),cut,sz(a),sz(b),
a[0]*b[0],a[sz(a)-1]*b[sz(b)-1],out,op);
}
template<engine E>
vector<typename E::value_type>middle_product(std::span<const typename E::value_type>a,transformed<E>&ta,
std::span<const typename E::value_type>b,transformed<E>&tb){
using T=typename E::value_type;
if(sz(a)==0||sz(b)==0)return{};
assert(sz(a)>=sz(b));
vector<T>r(size_t(sz(a)-sz(b)+1));
middle_product<E>(a,ta,b,tb,std::span<T>(r));
return r;
}
}
// src/fft/series_core.hpp
namespace ecnerwala::series{
template<fft::engine E,bool exact_>
struct span{
using T=typename E::value_type;
using engine_t=E;
static constexpr bool exact_v=exact_;
span()=default;
explicit span(std::span<const T>s_):s(s_){}
template<bool oe>requires(oe!=exact_)
explicit(oe<exact_)span(span<E,oe>o):s(o.coeffs()){}
int len()const{return sz(s);}
const T&operator[](int i)const{return s[size_t(i)];}
auto begin()const{return s.begin();}
auto end()const{return s.end();}
std::span<const T>coeffs()const{return s;}
span first(int n)const{
assert(n<=len());
return span(s.first(size_t(n)));
}
private:
std::span<const T>s;
};
template<fft::engine E,bool exact_>
struct vec:public std::vector<typename E::value_type>{
using T=typename E::value_type;
using engine_t=E;
static constexpr bool exact_v=exact_;
using std::vector<T>::vector;
operator span<E,exact_>()const{
return span<E,exact_>(std::span<const T>(*this));
}
template<bool oe>requires(oe!=exact_)
explicit(oe<exact_)vec(const vec<E,oe>&p):std::vector<T>(p){}
template<bool oe>requires(oe!=exact_)
explicit(oe<exact_)vec(vec<E,oe>&&p):std::vector<T>(std::move(p)){}
explicit vec(std::vector<T>v):std::vector<T>(std::move(v)){}
explicit vec(span<E,exact_>s):std::vector<T>(s.begin(),s.end()){}
explicit vec(span<E,!exact_>s):std::vector<T>(s.begin(),s.end()){}
span<E,exact_>first(int n)const{return span<E,exact_>(*this).first(n);}
int len()const{
return int(this->size());
}
int degree()const requires(exact_){
return len()-1;
}
void extend(int sz){
assert(sz>=len());
this->resize(sz);
}
void shrink(int sz){
assert(sz<=len());
this->resize(sz);
}
void shift_trunc(int n=1)requires(!exact_){
assert(n>=0&&n<=len());
std::rotate(this->begin(),this->end()-n,this->end());
std::fill(this->begin(),this->begin()+n,T(0));
}
void unshift_trunc(int n=1)requires(!exact_){
assert(n>=0&&n<=len());
std::fill(this->begin(),this->begin()+n,T(0));
std::rotate(this->begin(),this->begin()+n,this->end());
}
template<bool oe>requires(exact_<=oe)
vec&operator+=(const vec<E,oe>&o){
if constexpr(exact_){if(o.len()>len())this->resize(o.len());}
else if constexpr(!oe){if(o.len()<len())this->resize(o.len());}
for(int i=0;i<std::min(len(),o.len());i++){
(*this)[i]+=o[i];
}
return*this;
}
template<bool oe>requires(exact_<=oe)
vec&operator-=(const vec<E,oe>&o){
if constexpr(exact_){if(o.len()>len())this->resize(o.len());}
else if constexpr(!oe){if(o.len()<len())this->resize(o.len());}
for(int i=0;i<std::min(len(),o.len());i++){
(*this)[i]-=o[i];
}
return*this;
}
vec&operator*=(const T&n){
for(auto&v:*this)v*=n;
return*this;
}
friend vec operator*(const vec&a,const T&n){
vec r(a.size());
for(int i=0;i<a.len();i++){
r[i]=a[i]*n;
}
return r;
}
friend vec operator*(const T&n,const vec&a){
vec r(a.size());
for(int i=0;i<a.len();i++){
r[i]=n*a[i];
}
return r;
}
vec&operator*=(const vec&o){
return*this=(*this)*o;
}
};
template<fft::engine E>using exact=vec<E,true>;
template<fft::engine E>using trunc=vec<E,false>;
template<typename S>
concept like=fft::engine<typename S::engine_t>&&requires(const S&s,int i){
{S::exact_v}->std::convertible_to<bool>;
{s.len()}->std::same_as<int>;
{s[i]}->std::convertible_to<const typename S::engine_t::value_type&>;
{std::span<const typename S::engine_t::value_type>(s)};
requires std::convertible_to<const S&,span<typename S::engine_t,S::exact_v>>;
{s.first(i)}->std::convertible_to<span<typename S::engine_t,S::exact_v>>;
};
template<typename S>
concept exact_like=like<S>&&S::exact_v;
template<typename S>
concept trunc_like=like<S>&&!S::exact_v;
template<typename S>
concept has_cache=like<S>&&requires(const S&s){
{s.cache()}->std::same_as<fft::transformed<typename S::engine_t>&>;
};
template<fft::engine E,bool exact_>
struct maybe_cached;
template<fft::engine E,bool exact_>
struct cached_span{
using engine_t=E;
static constexpr bool exact_v=exact_;
span<E,exact_>s;
std::reference_wrapper<fft::transformed<E>>f;
cached_span(span<E,exact_>s_,fft::transformed<E>&f_):s(s_),f(f_){}
template<bool oe>requires(oe!=exact_)
explicit(oe<exact_)cached_span(cached_span<E,oe>o):s(span<E,exact_>(o.s)),f(o.f){}
int len()const{return s.len();}
const typename E::value_type&operator[](int i)const{return s[i];}
operator std::span<const typename E::value_type>()const{return s.coeffs();}
operator span<E,exact_>()const{return s;}
maybe_cached<E,exact_>first(int n)const;
fft::transformed<E>&cache()const{return f;}
};
template<typename S>
concept has_cache_opt=like<S>&&requires(const S&s){
{s.cache_opt()}->std::same_as<std::optional<std::reference_wrapper<fft::transformed<typename S::engine_t>>>>;
};
namespace detail{
template<like S>
std::optional<std::reference_wrapper<fft::transformed<typename S::engine_t>>>cache_of(const S&s){
if constexpr(has_cache<S>)return s.cache();
else if constexpr(has_cache_opt<S>)return s.cache_opt();
else return std::nullopt;
}
}
template<fft::engine E,bool exact_>
struct maybe_cached{
using T=typename E::value_type;
using engine_t=E;
static constexpr bool exact_v=exact_;
span<E,exact_>s;
std::optional<std::reference_wrapper<fft::transformed<E>>>f;
explicit maybe_cached(span<E,exact_>s_):s(s_){}
maybe_cached(span<E,exact_>s_,fft::transformed<E>&f_):s(s_),f(f_){}
maybe_cached(cached_span<E,exact_>c):s(c.s),f(c.f){}
template<like S>requires std::same_as<typename S::engine_t,E>&&(S::exact_v==exact_)
maybe_cached(const S&o):s(o),f(detail::cache_of(o)){}
int len()const{return s.len();}
const T&operator[](int i)const{return s[i];}
operator std::span<const T>()const{return s.coeffs();}
operator span<E,exact_>()const{return s;}
maybe_cached first(int n)const{
return n==len()?*this:maybe_cached(s.first(n));
}
std::optional<std::reference_wrapper<fft::transformed<E>>>cache_opt()const{return f;}
};
template<fft::engine E,bool exact_>
maybe_cached<E,exact_>cached_span<E,exact_>::first(int n)const{
return maybe_cached<E,exact_>(*this).first(n);
}
template<fft::engine E>
struct resized{
using T=typename E::value_type;
using engine_t=E;
static constexpr bool exact_v=false;
trunc<E>s;
std::optional<std::reference_wrapper<fft::transformed<E>>>f;
int len()const{return s.len();}
const T&operator[](int i)const{return s[size_t(i)];}
operator std::span<const T>()const{return std::span<const T>(s);}
operator span<E,false>()const{return s;}
maybe_cached<E,false>first(int n)const{
span<E,false>v=s;
if(n==len()&&f)return{v,f->get()};
return maybe_cached<E,false>(v.first(n));
}
std::optional<std::reference_wrapper<fft::transformed<E>>>cache_opt()const{return f;}
};
template<like S>
resized<typename S::engine_t>with_len(const S&s,int n){
using E=typename S::engine_t;
using T=typename E::value_type;
auto p=s.first(std::min(n,s.len()));
resized<E>r;
r.s.assign(size_t(n),T{});
std::span<const T>pc(p);
std::copy(pc.begin(),pc.end(),r.s.begin());
r.f=detail::cache_of(p);
return r;
}
template<typename S>
concept has_prefix_cache=like<S>&&!S::exact_v&&requires(const S&s,int n){
{s.prefix_cache(n)}->std::same_as<fft::transformed<typename S::engine_t>&>;
};
template<fft::engine E,bool exact_>
struct cached{
using T=typename E::value_type;
using engine_t=E;
static constexpr bool exact_v=exact_;
cached()=default;
cached(vec<E,exact_>&&s_):s(std::move(s_)){}
explicit cached(const vec<E,exact_>&s_):s(s_){}
operator vec<E,exact_>()&&{return std::move(s);}
int len()const{return s.len();}
const vec<E,exact_>&uncached()const{return s;}
const T&operator[](int i)const{return s[size_t(i)];}
auto begin()const{return s.cbegin();}
auto end()const{return s.cend();}
operator span<E,exact_>()const{return s;}
maybe_cached<E,exact_>first(int n)const{
return n==len()?maybe_cached<E,exact_>(s,f):maybe_cached<E,exact_>(s.first(n));
}
fft::transformed<E>&cache()const{return f;}
template<like S>
friend bool operator==(const cached&a,const S&b){
span<E,S::exact_v>bs=b;
return a.len()==bs.len()&&std::equal(a.s.begin(),a.s.end(),bs.begin());
}
private:
vec<E,exact_>s;
mutable fft::transformed<E>f;
};
template<fft::engine E>using cached_exact=cached<E,true>;
template<fft::engine E>using cached_trunc=cached<E,false>;
namespace detail{
template<like S>
cached_span<typename S::engine_t,S::exact_v>as_cached_span(const S&s,fft::transformed<typename S::engine_t>&tmp){
auto co=cache_of(s);
return{s,co?co->get():tmp};
}
}
template<trunc_like S>
trunc<typename S::engine_t>ps_inv(const S&a){
using E=typename S::engine_t;
using T=typename E::value_type;
int N=a.len();
trunc<E>r(size_t(N),T{});
if(N==0)return r;
int s=nextPow2(N);
std::vector<T>b(size_t(s),T{});
b[0]=inv(a[0]);
for(int n=1;n<N;n*=2){
int m=2*n;
auto ta=E::transform(a.first(std::min(N,m)),m);
auto tb=E::transform(std::span<const T>(b).first(n),m);
auto e=fft::buffer_pool<T>::get(m);
E::finish(E::mul(ta,tb,m),e.span());
for(int i=0;i<n;i++)e[i]=T{};
auto te=E::transform(std::span<const T>(e.span()),m);
auto c=fft::buffer_pool<T>::get(m);
E::finish(E::mul(tb,te,m),c.span());
for(int i=n;i<std::min(m,N);i++)b[i]=-c[i];
}
std::copy(b.begin(),b.begin()+N,r.begin());
return r;
}
template<like A>
auto square(const A&a){
using E=typename A::engine_t;
using T=typename E::value_type;
fft::transformed<E>ta_;
auto av=detail::as_cached_span(a,ta_);
if constexpr(A::exact_v){
std::vector<T>coeffs;
fft::transformed<E>f;
fft::square_cached<E>(av,av.cache(),coeffs,f);
cached<E,true>w(exact<E>(std::move(coeffs)));
w.cache()=std::move(f);
return w;
}else{
trunc<E>r(size_t(a.len()),T{});
fft::square<E>(av,av.cache(),std::span<T>(r));
return r;
}
}
template<like A,like B,like C,like D>
requires fft::same_engine<A,B>&&fft::same_engine<A,C>&&fft::same_engine<A,D>
&&A::exact_v&&B::exact_v&&C::exact_v&&D::exact_v
cached<typename A::engine_t,true>multiply_add2(
const A&a,const B&b,const C&c,const D&d){
using E=typename A::engine_t;
using T=typename E::value_type;
fft::transformed<E>ta_,tb_,tc_,td_;
auto av=detail::as_cached_span(a,ta_),bv=detail::as_cached_span(b,tb_);
auto cv=detail::as_cached_span(c,tc_),dv=detail::as_cached_span(d,td_);
std::vector<T>coeffs;
fft::transformed<E>f;
fft::multiply_add2_cached<E>(
av,av.cache(),
bv,bv.cache(),
cv,cv.cache(),
dv,dv.cache(),
coeffs,f
);
cached<E,true>w(exact<E>(std::move(coeffs)));
w.cache()=std::move(f);
return w;
}
template<like A,exact_like B>requires fft::same_engine<A,B>
vec<typename A::engine_t,A::exact_v>middle_product(const A&a,const B&b){
using E=typename A::engine_t;
fft::transformed<E>ta_,tb_;
auto av=detail::as_cached_span(a,ta_);
auto bv=detail::as_cached_span(b,tb_);
return vec<E,A::exact_v>(fft::middle_product<E>(
av,av.cache(),
bv,bv.cache()
));
}
namespace detail{
template<bool ea,bool eb>int product_prec(int la,int lb){
if constexpr(ea&&eb)return la>0&&lb>0?la+lb-1:0;
else return ea?lb:eb?la:std::min(la,lb);
}
template<like S>
auto product_operand(const S&s,int prec,fft::transformed<typename S::engine_t>&tmp){
using E=typename S::engine_t;
if constexpr(has_prefix_cache<S>){
return s.first(std::min(s.len(),nextPow2(prec)));
}else{
span<E,S::exact_v>v=s;
int used=std::min(v.len(),prec);
if(auto co=cache_of(s)){
if(s.len()<=prec||fft::detail::conv_size_for(s.len()+prec-1).n
==fft::detail::conv_size_for(2*prec-1).n){
return cached_span<E,S::exact_v>{v,co->get()};
}
}
return cached_span<E,S::exact_v>{v.first(used),tmp};
}
}
}
template<like A,like B>requires fft::same_engine<A,B>
vec<typename A::engine_t,A::exact_v&&B::exact_v>operator+(const A&a,const B&b){
using T=typename A::engine_t::value_type;
int n=(A::exact_v&&B::exact_v)?std::max(a.len(),b.len())
:A::exact_v?b.len():B::exact_v?a.len():std::min(a.len(),b.len());
vec<typename A::engine_t,A::exact_v&&B::exact_v>r(size_t(n),T(0));
for(int i=0;i<n;i++){
r[i]=(i<a.len()?a[i]:T(0))+(i<b.len()?b[i]:T(0));
}
return r;
}
template<like A,like B>requires fft::same_engine<A,B>
vec<typename A::engine_t,A::exact_v&&B::exact_v>operator-(const A&a,const B&b){
using T=typename A::engine_t::value_type;
int n=(A::exact_v&&B::exact_v)?std::max(a.len(),b.len())
:A::exact_v?b.len():B::exact_v?a.len():std::min(a.len(),b.len());
vec<typename A::engine_t,A::exact_v&&B::exact_v>r(size_t(n),T(0));
for(int i=0;i<n;i++){
r[i]=(i<a.len()?a[i]:T(0))-(i<b.len()?b[i]:T(0));
}
return r;
}
template<like A,like B>requires fft::same_engine<A,B>
auto operator*(const A&a,const B&b){
using E=typename A::engine_t;
using T=typename E::value_type;
constexpr bool ea=A::exact_v,eb=B::exact_v;
int prec=detail::product_prec<ea,eb>(a.len(),b.len());
if(prec==0||a.len()==0||b.len()==0){
if constexpr(ea&&eb)return cached<E,true>{};
else return trunc<E>(size_t(prec),T(0));
}
fft::transformed<E>ta_,tb_;
auto va=detail::product_operand(a,prec,ta_);
auto vb=detail::product_operand(b,prec,tb_);
if constexpr(ea&&eb){
std::vector<T>coeffs;
fft::transformed<E>f;
auto ca=detail::as_cached_span(va,ta_),cb=detail::as_cached_span(vb,tb_);
fft::multiply_cached<E>(
ca,ca.cache(),
cb,cb.cache(),
coeffs,f
);
cached<E,true>w(exact<E>(std::move(coeffs)));
w.cache()=std::move(f);
return w;
}else{
trunc<E>r(size_t(prec),T(0));
auto ca=detail::as_cached_span(va,ta_),cb=detail::as_cached_span(vb,tb_);
fft::multiply<E>(
ca,ca.cache(),
cb,cb.cache(),
std::span<T>(r)
);
return r;
}
}
template<fft::engine E>
struct prefix_cached{
using T=typename E::value_type;
using engine_t=E;
static constexpr bool exact_v=false;
prefix_cached()=default;
prefix_cached(trunc<E>&&s_):s(std::move(s_)){}
explicit prefix_cached(const trunc<E>&s_):s(s_){}
operator trunc<E>()&&{return std::move(s);}
int len()const{return s.len();}
const trunc<E>&uncached()const{return s;}
const T&operator[](int i)const{return s[size_t(i)];}
auto begin()const{return s.cbegin();}
auto end()const{return s.cend();}
operator span<E,false>()const{return s;}
void append(std::span<const T>tail){
s.insert(s.end(),tail.begin(),tail.end());
}
maybe_cached<E,false>first(int k)const{
assert(k<=len());
span<E,false>v=s.first(k);
int n=nextPow2(k);
if(std::min(n,len())==k)return{v,prefix_cache(n)};
return maybe_cached<E,false>(v);
}
fft::transformed<E>&cache()const{return prefix_cache(nextPow2(len()));}
fft::transformed<E>&prefix_cache(int n)const{
assert(n>0&&!(n&(n-1)));
int k=__builtin_ctz(unsigned(n));
if(k>=sz(caches))caches.resize(size_t(k)+1);
auto&c=caches[k];
int e=std::min(n,len());
if(c.len!=e){
c.t=E::transform(s.first(e),2*n);
c.len=e;
}
return c.t;
}
private:
trunc<E>s;
struct entry{fft::transformed<E>t;int len=0;};
mutable std::vector<entry>caches;
};
}
// src/fft/poly.hpp
namespace ecnerwala::poly{
template<fft::engine E>struct vec{
using T=typename E::value_type;
using engine_t=E;
series::exact<E>c;
vec()=default;
explicit vec(int len):c(size_t(len),T{}){}
vec(std::initializer_list<T>coeffs):c(std::rbegin(coeffs),std::rend(coeffs)){}
explicit vec(std::span<const T>coeffs):c(coeffs.rbegin(),coeffs.rend()){}
const series::exact<E>&rev_series()const{return c;}
static vec from_rev_series(series::exact<E>s){
vec r;
r.c=std::move(s);
return r;
}
series::trunc<E>unrev_series(int n)const{
series::trunc<E>r(size_t(n),T{});
std::copy(begin(),begin()+std::min(n,len()),r.begin());
return r;
}
auto begin(){return c.rbegin();}
auto end(){return c.rend();}
auto begin()const{return c.rbegin();}
auto end()const{return c.rend();}
int len()const{return c.len();}
int degree()const{return len()-1;}
T&operator[](int i){return c[len()-1-i];}
const T&operator[](int i)const{return c[len()-1-i];}
T leading()const{return c.front();}
void shift(int k=1){
if(len()>0)c.insert(c.end(),size_t(k),T(0));
}
void resize(int n){
if(n>=len())c.insert(c.begin(),size_t(n-len()),T(0));
else c.erase(c.begin(),c.begin()+(len()-n));
}
T operator()(const T&x)const{
T r{};
for(const T&v:c)r=r*x+v;
return r;
}
vec&operator+=(const vec&o){
if(o.len()>len())resize(o.len());
for(int i=0;i<o.len();i++)(*this)[i]+=o[i];
return*this;
}
friend vec operator+(vec a,const vec&b){a+=b;return a;}
vec&operator-=(const vec&o){
if(o.len()>len())resize(o.len());
for(int i=0;i<o.len();i++)(*this)[i]-=o[i];
return*this;
}
friend vec operator-(vec a,const vec&b){a-=b;return a;}
friend bool operator==(const vec&a,const vec&b){return a.c==b.c;}
vec&operator*=(const T&n){for(T&v:c)v*=n;return*this;}
friend vec operator*(vec a,const T&n){a*=n;return a;}
friend vec operator*(const T&n,vec a){a*=n;return a;}
vec&operator*=(const vec&o){return*this=(*this)*o;}
};
template<typename P>
concept like=requires(const P&p){
typename P::engine_t;
{p.len()}->std::same_as<int>;
p.rev_series();
requires series::like<std::remove_cvref_t<decltype(p.rev_series())>>;
requires std::remove_cvref_t<decltype(p.rev_series())>::exact_v;
};
template<fft::engine E>
struct cached{
using T=typename E::value_type;
using engine_t=E;
cached()=default;
cached(vec<E>&&p):c(std::move(p.c)){}
explicit cached(const vec<E>&p):c(p.c){}
operator vec<E>()&&{return vec<E>::from_rev_series(std::move(c));}
const series::cached_exact<E>&rev_series()const{return c;}
static cached from_rev_series(series::cached_exact<E>s){
cached r;
r.c=std::move(s);
return r;
}
int len()const{return c.len();}
int degree()const{return len()-1;}
const T&operator[](int i)const{return c[len()-1-i];}
T leading()const{return c[0];}
T operator()(const T&x)const{
T r{};
for(const T&v:c)r=r*x+v;
return r;
}
private:
series::cached_exact<E>c;
};
template<like A,like B>requires fft::same_engine<A,B>
cached<typename A::engine_t>operator*(const A&a,const B&b){
return cached<typename A::engine_t>::from_rev_series(a.rev_series()*b.rev_series());
}
template<like A>
cached<typename A::engine_t>square(const A&a){
return cached<typename A::engine_t>::from_rev_series(square(a.rev_series()));
}
template<like A,like B,like C,like D>
requires fft::same_engine<A,B>&&fft::same_engine<A,C>&&fft::same_engine<A,D>
cached<typename A::engine_t>multiply_add2(
const A&a,const B&b,const C&c,const D&d){
return cached<typename A::engine_t>::from_rev_series(
multiply_add2(a.rev_series(),b.rev_series(),c.rev_series(),d.rev_series()));
}
template<like A,like B>requires fft::same_engine<A,B>
bool operator==(const A&a,const B&b){
return a.rev_series()==b.rev_series();
}
template<fft::engine E>
struct form{
using T=typename E::value_type;
series::cached_exact<E>c;
form()=default;
explicit form(int len):c(series::exact<E>(size_t(len),T{})){}
const series::cached_exact<E>&rev_series()const{return c;}
static form from_rev_series(series::cached_exact<E>s){
form r;
r.c=std::move(s);
return r;
}
static form from_poly(const vec<E>&p){return from_rev_series(series::cached_exact<E>(p.rev_series()));}
int len()const{return c.len();}
form for_length(int n)const{
auto r=series::exact<E>(c);
if(n>=len())r.insert(r.begin(),size_t(n-len()),T(0));
else r.erase(r.begin(),r.begin()+(len()-n));
return from_rev_series(std::move(r));
}
static form polynomial_evaluation(T z,int len){
series::exact<E>k(size_t(len),T{});
T p=T(1);
for(int i=0;i<len;i++){k[i]=p;p*=z;}
return from_rev_series(std::move(k));
}
template<like P>
T operator()(const P&p)const{
assert(p.len()<=len());
T r{};
for(int i=0;i<p.len();i++)r+=c[i]*p[i];
return r;
}
template<like P>
form composed_with(const P&q)const{
assert(q.len()>0&&q.len()<=len());
return from_rev_series(middle_product(c,q.rev_series()));
}
template<series::like S>requires std::same_as<typename S::engine_t,E>
form composed_with(const S&s)const{
if constexpr(!S::exact_v)assert(s.len()>=len());
series::vec<E,S::exact_v>r=c*s;
r.resize(size_t(len()));
return from_rev_series(series::exact<E>(std::move(r)));
}
};
template<fft::engine E>
struct subproduct_tree{
using T=typename E::value_type;
int N;
std::vector<cached<E>>nodes;
explicit subproduct_tree(std::span<const T>pts):N(sz(pts)),nodes(size_t(2)*N){
assert(N>0);
for(int i=0;i<N;i++){
nodes[N+i]=vec<E>{-pts[i],T(1)};
}
for(int i=N-1;i>0;i--){
nodes[i]=nodes[2*i]*nodes[2*i+1];
}
}
int size(int i)const{return nodes[i].len()-1;}
const cached<E>&prod(int i)const{return nodes[i];}
std::vector<T>pushdown(form<E>f)const{
assert(f.len()==N);
std::vector<form<E>>down(size_t(2)*N);
down[1]=std::move(f);
for(int i=1;i<N;i++){
down[2*i+0]=down[i].composed_with(nodes[2*i+1]);
down[2*i+1]=down[i].composed_with(nodes[2*i+0]);
down[i]=form<E>{};
}
std::vector<T>out(size_t(N),T{});
for(int i=0;i<N;i++)out[i]=down[N+i].rev_series()[0];
return out;
}
cached<E>combine_up(std::span<const T>leaf_vals)const{
assert(sz(leaf_vals)==N);
std::vector<cached<E>>up(size_t(2)*N);
for(int i=0;i<N;i++){
up[N+i]=vec<E>{leaf_vals[i]};
}
for(int i=N-1;i>0;i--){
up[i]=multiply_add2(up[2*i+0],nodes[2*i+1],up[2*i+1],nodes[2*i+0]);
up[2*i+0]=cached<E>{};
up[2*i+1]=cached<E>{};
}
return std::move(up[1]);
}
};
template<fft::engine E>
std::vector<typename E::value_type>multipoint(
const vec<E>&p,
std::span<const typename E::value_type>pts
){
if(pts.empty())return{};
int N=sz(pts);
subproduct_tree<E>tree{pts};
auto q=series::trunc<E>(tree.prod(1).rev_series());
q.resize(p.len());
form<E>f=form<E>::from_poly(p).composed_with(ps_inv(q));
return tree.pushdown(f.for_length(N));
}
template<fft::engine E>
vec<E>interpolate(
std::span<const typename E::value_type>pts,
std::span<const typename E::value_type>vals
){
using T=typename E::value_type;
assert(sz(pts)==sz(vals));
if(pts.empty())return{};
int N=sz(pts);
using ps=series::trunc<E>;
subproduct_tree<E>tree{pts};
auto root=ps(tree.prod(1).rev_series());
root.shrink(N);
ps deriv_root=root;
for(int i=0;i<N;i++){
deriv_root[i]*=T(N-i);
}
std::vector<T>denoms=tree.pushdown(
form<E>::from_rev_series(series::exact<E>(ps_inv(root)*deriv_root))
);
std::vector<T>leaf_vals(size_t(N),T{});
for(int i=0;i<N;i++)leaf_vals[i]=vals[i]/denoms[i];
return tree.combine_up(std::span<const T>(leaf_vals));
}
}
// src/fft/series.hpp
namespace ecnerwala::series{
template<like S>
vec<typename S::engine_t,S::exact_v>stretch(const S&a_,int n){
using E=typename S::engine_t;
span<E,S::exact_v>a=a_;
vec<E,S::exact_v>r(size_t(a.len()));
for(int i=0;i*n<a.len();i++){
r[i*n]=a[i];
}
return r;
}
template<like S>
vec<typename S::engine_t,S::exact_v>deriv_shift(const S&a_){
using E=typename S::engine_t;
span<E,S::exact_v>a=a_;
vec<E,S::exact_v>r(a.begin(),a.end());
for(int i=0;i<r.len();i++){
r[i]*=i;
}
return r;
}
template<like S>
vec<typename S::engine_t,S::exact_v>integ_shift(const S&a_){
using E=typename S::engine_t;
using T=typename E::value_type;
span<E,S::exact_v>a=a_;
assert(a[0]==0);
vec<E,S::exact_v>r(a.begin(),a.end());
T f=1;
for(int i=1;i<r.len();i++){
r[i]*=f;
f*=i;
}
f=inv(f);
for(int i=r.len()-1;i>0;i--){
r[i]*=f;
f*=i;
}
return r;
}
template<like S>
vec<typename S::engine_t,S::exact_v>integ_shift_offset(const S&a_,int offset){
using E=typename S::engine_t;
using T=typename E::value_type;
span<E,S::exact_v>a=a_;
vec<E,S::exact_v>r(a.begin(),a.end());
T f=1;
for(int i=0;i<r.len();i++){
r[i]*=f;
f*=i+offset;
}
assert(f!=0);
f=inv(f);
for(int i=r.len()-1;i>=0;i--){
r[i]*=f;
f*=i+offset;
}
return r;
}
template<trunc_like S>
trunc<typename S::engine_t>deriv_shift_log(const S&a){
return deriv_shift(a)*ps_inv(a);
}
template<trunc_like S>
trunc<typename S::engine_t>ps_log(const S&a){
assert(a[0]==1);
return integ_shift(deriv_shift_log(a));
}
template<trunc_like S>
trunc<typename S::engine_t>ps_exp(const S&a_){
using E=typename S::engine_t;
using T=typename E::value_type;
span<E,false>a=a_;
assert(a.len()>=1);
assert(a[0]==0);
trunc<E>r(1,T(1));r.reserve(size_t(a.len()));
trunc<E>invR(1,T(1));invR.reserve(size_t(a.len()));
while(r.len()<a.len()){
int o_sz=r.len();
int n_sz=std::min(o_sz*2,a.len());
trunc<E>t=deriv_shift(trunc<E>(a.begin(),a.begin()+o_sz));
fft::multiply_circular<E>(std::span<const T>(t),std::span<const T>(r).first(o_sz),std::span<T>(t),o_sz);
t=deriv_shift(r)-t;
t*=invR;
t.resize(size_t(n_sz-o_sz));
trunc<E>v(a.begin()+o_sz,a.begin()+n_sz);
v-=integ_shift_offset(t,o_sz);
v*=r;
r.resize(size_t(n_sz));
std::copy(v.begin(),v.end(),r.begin()+o_sz);
if(r.len()<a.len()){
assert(r.len()==2*invR.len());
int n=invR.len();
int nn=r.len();
trunc<E>tmp(size_t(4)*n);
fft::square<E>(std::span<const T>(invR).first(n),std::span<T>(tmp));
fft::multiply<E>(std::span<const T>(tmp).first(nn),std::span<const T>(r).first(nn),std::span<T>(tmp));
invR.resize(size_t(nn));
for(int i=n;i<nn;i++)invR[i]=-tmp[i];
}
}
return r;
}
template<trunc_like S>
trunc<typename S::engine_t>ps_pow_monic(const S&a_,typename S::engine_t::value_type k){
using E=typename S::engine_t;
span<E,false>a=a_;
if(a.len()==0)return{};
assert(a[0]==1);
trunc<E>l=ps_log(a_);
l*=k;
return ps_exp(l);
}
template<trunc_like S>
trunc<typename S::engine_t>ps_pow(const S&a_,int64_t k){
using E=typename S::engine_t;
using T=typename E::value_type;
span<E,false>a=a_;
assert(k>=0);
if(k==0){
trunc<E>r(size_t(a.len()),T(0));
if(r.len()>0)r[0]=T(1);
return r;
}
int st=0;
while(st<a.len()&&a[st]==0)st++;
if(st>0&&k>(a.len()-1)/st){
return trunc<E>(size_t(a.len()),T(0));
}
trunc<E>r(a.begin()+st,a.end()-(st*(k-1)));
T leading_coeff=r[0];
r*=inv(leading_coeff);
r=ps_pow_monic(r,T(k));
r*=power(leading_coeff,k);
r.insert(r.begin(),size_t(st*k),T(0));
assert(r.len()==a.len());
return r;
}
template<trunc_like S>
trunc<typename S::engine_t>to_newton_sums(const S&a,int deg){
auto r=deriv_shift_log(a);
r[0]=deg;
for(int i=1;i<r.len();i++)r[i]=-r[i];
return r;
}
template<trunc_like S>
trunc<typename S::engine_t>from_newton_sums(const S&s_,int deg){
using E=typename S::engine_t;
span<E,false>s=s_;
assert(s[0]==deg);
trunc<E>r(s.begin(),s.end());
r[0]=0;
for(int i=1;i<r.len();i++)r[i]=-r[i];
return ps_exp(integ_shift(std::move(r)));
}
template<trunc_like S>
trunc<typename S::engine_t>euler_transform(const S&a){
using E=typename S::engine_t;
trunc<E>r=deriv_shift(a);
std::vector<bool>is_prime(size_t(r.len()),true);
for(int p=2;p<r.len();p++){
if(!is_prime[p])continue;
for(int i=1;i*p<r.len();i++){
r[i*p]+=r[i];
is_prime[i*p]=false;
}
}
return ps_exp(integ_shift(r));
}
template<trunc_like S>
trunc<typename S::engine_t>inverse_euler_transform(const S&a){
using E=typename S::engine_t;
trunc<E>r=deriv_shift(ps_log(a));
std::vector<bool>is_prime(size_t(r.len()),true);
for(int p=2;p<r.len();p++){
if(!is_prime[p])continue;
for(int i=(r.len()-1)/p;i>=1;i--){
r[i*p]-=r[i];
is_prime[i*p]=false;
}
}
return integ_shift(r);
}
template<fft::engine E>struct packed_bivariate{
using T=typename E::value_type;
int L,l;
std::vector<T>c;
packed_bivariate(int L_,std::span<const T>g):L(L_),l(0),c(size_t(4)<<L){
c[0]=T(1);
for(int i=0;i<sz(g);i++)c[(2<<L)+i]=-g[i];
}
fft::transformed<E>advance(){
int B=4<<L;
auto tq=E::transform(std::span<const T>(c),B);
auto tn=E::negate_arg(tq,B);
E::finish(
E::downsample(E::mul(tq,tn,B),B/2,false),
std::span<T>(c).first(B/2)
);
l++;
for(int i=0;i<(2<<(L-l));i++){
c[(2<<L)+i]=c[i];
c[i]=T(0);
}
c[2<<L]-=T(1);
c[0]=T(1);
std::fill(c.begin()+(2<<L)+(1<<(L-l)),c.end(),T(0));
for(int i=0;i<(2<<L);i+=2<<(L-l)){
for(int j=0;j<(1<<(L-l));j++){
c[i+(1<<(L-l))+j]=T(0);
}
}
return tn;
}
};
template<trunc_like SF,trunc_like SG>requires fft::same_engine<SF,SG>
trunc<typename SF::engine_t>ps_compose(const SF&f_,const SG&g_){
using E=typename SF::engine_t;
using T=typename E::value_type;
span<E,false>f=f_;
span<E,false>g=g_;
if(g.len()==0)return{};
int m=f.len();
int n=g.len();
int L=__builtin_ctz(unsigned(nextPow2(n)));
int B=4<<L;
packed_bivariate<E>Q(L,g.coeffs());
std::vector<fft::transformed<E>>tneg;
tneg.reserve(L);
for(int l=1;l<=L;l++)tneg.push_back(Q.advance());
trunc<E>P;
{
P=trunc<E>(f.begin(),f.end());
std::reverse(P.begin(),P.end());
trunc<E>QL((1<<L)+1);
for(int i=0;i<=(1<<L);i++){
QL[i]=Q.c[2*i];
}
QL.resize(size_t(m),T(0));
P*=ps_inv(QL);
std::reverse(P.begin(),P.end());
P.resize(size_t(1)<<L,T(0));
std::reverse(P.begin(),P.end());
P.resize(size_t(B),T(0));
for(int i=(1<<L)-1;i>0;i--){
P[2*i]=P[i];
P[i]=T(0);
}
}
for(int l=L-1;l>=0;l--){
for(int i=(2<<L)-1;i>0;i--){
T v=P[i];
P[2*i]=((2*i)&(1<<(L-l)))?T(0):v;
P[i]=T(0);
}
auto tp=E::transform(std::span<const T>(P),B);
E::finish(E::mul(tneg[l],tp,B),std::span<T>(P));
for(int i=0;i<(2<<L);i++){
P[i]=P[(2<<L)+i];
P[(2<<L)+i]=T(0);
}
}
return trunc<E>(P.begin(),P.begin()+n);
}
template<exact_like P,exact_like Q>requires fft::same_engine<P,Q>
P::engine_t::value_type kth_term_of_rational_function(
const P&p,
const Q&q,
uint64_t k
){
using E=P::engine_t;
using T=E::value_type;
assert(q.len()>0&&q[0]!=T(0));
if(p.len()==0)return T(0);
int d=std::max(p.len()+1,q.len());
assert(d>=2);
int n=nextPow2((d-1)+d-1);
fft::transformed<E>tq,tp;
if(auto cq=detail::cache_of(q)){E::extend_to(cq->get(),n,q);tq=cq->get();}
if(auto cp=detail::cache_of(p)){E::extend_to(cp->get(),n,p);tp=cp->get();}
std::vector<T>p_buf(d-1,T(0));
std::ranges::copy(std::span<const T>(p),p_buf.begin());
std::vector<T>q_buf(d,T(0));
std::ranges::copy(std::span<const T>(q),q_buf.begin());
while(k>0){
E::extend_to(tq,n,q_buf);
auto tnq=E::negate_arg(tq,n);
E::extend_to(tp,n,p_buf);
auto ntp=E::downsample(E::mul(tp,tnq,n),n/2,bool(k&1));
assert(ntp.size()==n/2);
if constexpr(std::same_as<typename E::product,typename E::transformed>){
tp=ntp;
}else{
tp={};
}
E::finish(std::move(ntp),std::span(p_buf));
k>>=1;
if(!k){
q_buf[0]*=q_buf[0];
break;
}
auto ntq=E::downsample(E::mul(tq,tnq,n),n/2,false);
assert(ntq.size()==n/2);
if constexpr(std::same_as<typename E::product,typename E::transformed>){
tq=ntq;
}else{
tq={};
}
if(n/2==d-1){
T v0=q_buf[0]*q_buf[0];
E::finish(std::move(ntq),std::span(q_buf).first(d-1));
q_buf[d-1]=std::exchange(q_buf[0],v0)-v0;
}else{
E::finish(std::move(ntq),std::span(q_buf));
}
}
return p_buf[0]*inv(q_buf[0]);
}
template<trunc_like S,exact_like Q>requires fft::same_engine<S,Q>
S::engine_t::value_type kth_term_of_linear_recurrence(
const S&s,
const Q&q,
uint64_t k
){
using E=S::engine_t;
using T=E::value_type;
assert(q.len()>0&&q[0]!=T(0));
assert(s.len()>=q.len()-1);
fft::transformed<E>tq;
auto q_cached=detail::as_cached_span(q,tq);
span<E,false>sv=s;
auto p=exact<E>(sv.first(q.len()-1)*q_cached);
return kth_term_of_rational_function(p,q_cached,k);
}
}
// src/modnum.hpp
template<typename T>T mod_inv_in_range(T a,T m){
T x=a,y=m;
T vx=1,vy=0;
bool swap=false;
while(x){
T k=y/x;
y%=x;
vy+=k*vx;
std::swap(x,y);
std::swap(vx,vy);
swap^=1;
}
assert(y==1);
return swap?vy:m-vy;
}
template<typename T>struct extended_gcd_result{
T gcd;
T coeff_a,coeff_b;
};
template<typename T>extended_gcd_result<T>extended_gcd(T a,T b){
T x=a,y=b;
T ax=1,ay=0;
T bx=0,by=1;
while(x){
T k=y/x;
y%=x;
ay-=k*ax;
by-=k*bx;
std::swap(x,y);
std::swap(ax,ay);
std::swap(bx,by);
}
return{y,ay,by};
}
template<typename T>T mod_inv(T a,T m){
a%=m;
a=a<0?a+m:a;
return mod_inv_in_range(a,m);
}
template<typename Self>
struct num_ops{
Self operator+()const{return static_cast<const Self&>(*this);}
Self operator-()const{return static_cast<const Self&>(*this).neg();}
friend Self operator++(Self&a,int){Self r=a;++a;return r;}
friend Self operator--(Self&a,int){Self r=a;--a;return r;}
friend Self operator+(const Self&a,const Self&b){return Self(a)+=b;}
friend Self operator-(const Self&a,const Self&b){return Self(a)-=b;}
friend Self operator*(const Self&a,const Self&b){return Self(a)*=b;}
friend Self operator/(const Self&a,const Self&b){return Self(a)/=b;}
friend bool operator!=(const Self&a,const Self&b){return!(a==b);}
friend Self neg(const Self&a){return a.neg();}
friend Self inv(const Self&a){return a.inv();}
};
template<typename Self,typename V>
struct mod_ops:num_ops<Self>{
static_assert(std::unsigned_integral<V>);
V v;
struct is_reduced_tag{};
mod_ops():v(0){}
mod_ops(V v_,is_reduced_tag):v(v_){assert(v<Self::MOD);}
template<std::integral I>mod_ops(I x):v(Self::reduce(x)){}
static Self from_reduced(V v){return Self(v,is_reduced_tag{});}
static V reduce(std::signed_integral auto x){
using U=std::make_unsigned_t<decltype(x)>;
return x<0?V(Self::MOD-1-Self::reduce(U(~x))):Self::reduce(U(x));
}
explicit operator V()const{return v;}
std::make_signed_t<V>balanced()const{
return std::make_signed_t<V>(Self::MOD-v>v?v:v-Self::MOD);
}
friend bool operator==(const Self&a,const Self&b){return a.v==b.v;}
friend std::ostream&operator<<(std::ostream&out,const Self&n){return out<<n.v;}
friend std::istream&operator>>(std::istream&in,Self&n){int64_t v_;in>>v_;n=Self(v_);return in;}
Self&operator++(){
++v;
if(v==Self::MOD)v=0;
return self();
}
Self&operator--(){
if(v==0)v=Self::MOD;
--v;
return self();
}
Self&operator+=(const Self&o){v=Self::sub_mod_raw(v,Self::MOD-o.v);return self();}
Self&operator-=(const Self&o){v=Self::sub_mod_raw(v,o.v);return self();}
Self&operator/=(const Self&o){return self()*=o.inv();}
static V sub_mod_raw(V a,V b){return a<b?a-b+Self::MOD:a-b;}
Self neg()const{return from_reduced(v?Self::MOD-v:0);}
Self inv()const{return from_reduced(mod_inv_in_range(v,Self::MOD));}
private:
Self&self(){return static_cast<Self&>(*this);}
};
template<auto MOD_>struct modnum:mod_ops<modnum<MOD_>,std::make_unsigned_t<decltype(MOD_)>>{
using Self=modnum;
static_assert(MOD_>0,"MOD must be positive");
using V=std::make_unsigned_t<decltype(MOD_)>;
static constexpr V MOD=V(MOD_);
using base=mod_ops<modnum,V>;
using base::base;
using base::v;
using base::reduce;
static V reduce(std::unsigned_integral auto x){return V(x%MOD);}
explicit operator std::make_signed_t<V>()const
requires(MOD<=V(std::numeric_limits<std::make_signed_t<V>>::max()))
{
return std::make_signed_t<V>(v);
}
Self&operator*=(const Self&o){
if constexpr(sizeof(V)<=4)v=V(uint64_t(v)*o.v%MOD);
else v=V(__uint128_t(v)*o.v%MOD);
return*this;
}
};
struct mod_goldilocks:mod_ops<mod_goldilocks,uint64_t>{
using Self=mod_goldilocks;
static constexpr uint64_t MOD=0xffffffff00000001ull;
static constexpr uint64_t EPS=-MOD;
static constexpr uint64_t PRIMITIVE_ROOT=2717;
using base=mod_ops<mod_goldilocks,uint64_t>;
using base::base;
using base::reduce;
mod_goldilocks()=default;
mod_goldilocks(__int128_t a):base(a<0?uint64_t(MOD-1-__uint128_t(~a)%MOD):uint64_t(__uint128_t(a)%MOD),is_reduced_tag{}){}
mod_goldilocks(__uint128_t a):base(uint64_t(a%MOD),is_reduced_tag{}){}
static uint64_t reduce(std::unsigned_integral auto x){
static_assert(sizeof(x)<=8);
uint64_t a=x;
return a>=MOD?a-MOD:a;
}
static uint64_t sub_mod_raw(uint64_t a,uint64_t b){
#if defined(__x86_64__)
uint64_t res_wrapped=a;
uint64_t adjustment=b;
asm(
"sub %[y], %[x]\n\t"
"sbb %k[y], %k[y]\n\t"
:[x]"+r"(res_wrapped),
[y]"+r"(adjustment)
:
:"cc"
);
#else
uint64_t res_wrapped=a-b;
uint64_t adjustment=(res_wrapped>a)?EPS:0;
#endif
return res_wrapped-adjustment;
}
static uint64_t reduce_u160_raw(uint64_t lo,uint32_t mi,uint64_t hi){
return sub_mod_raw(sub_mod_raw(lo,hi),MOD-(uint64_t(mi)<<32)+mi);
}
static uint64_t reduce_u128_raw(__uint128_t v){
uint64_t hi=uint64_t(v>>64);
uint64_t lo=uint64_t(v);
uint32_t hi_hi=uint32_t(hi>>32);
uint32_t hi_lo=uint32_t(hi);
return reduce_u160_raw(lo,hi_lo,hi_hi);
}
Self&operator*=(Self o){
v=reduce_u128_raw(__uint128_t(v)*__uint128_t(o.v));
return*this;
}
};
template<typename T>T power(T a,long long b){
assert(b>=0);
T r=1;while(b){if(b&1)r*=a;b>>=1;a*=a;}return r;
}
template<typename U,typename V>struct pairnum:num_ops<pairnum<U,V>>{
using Self=pairnum;
U u;
V v;
pairnum():u(0),v(0){}
pairnum(long long val):u(val),v(val){}
pairnum(const U&u_,const V&v_):u(u_),v(v_){}
friend std::ostream&operator<<(std::ostream&out,const Self&n){return out<<'('<<n.u<<','<<' '<<n.v<<')';}
friend std::istream&operator>>(std::istream&in,Self&n){long long val;in>>val;n=Self(val);return in;}
friend bool operator==(const Self&a,const Self&b){return a.u==b.u&&a.v==b.v;}
Self inv()const{
return Self(u.inv(),v.inv());
}
Self neg()const{
return Self(u.neg(),v.neg());
}
Self&operator++(){
++u,++v;
return*this;
}
Self&operator--(){
--u,--v;
return*this;
}
Self&operator+=(const Self&o){
u+=o.u;
v+=o.v;
return*this;
}
Self&operator-=(const Self&o){
u-=o.u;
v-=o.v;
return*this;
}
Self&operator*=(const Self&o){
u*=o.u;
v*=o.v;
return*this;
}
Self&operator/=(const Self&o){
u/=o.u;
v/=o.v;
return*this;
}
};
template<typename tag>struct dynamic_modnum:mod_ops<dynamic_modnum<tag>,uint32_t>{
using Self=dynamic_modnum;
private:
inline static uint32_t MOD_=0;
inline static uint64_t BARRETT_M=0;
public:
static constexpr uint32_t const&MOD=MOD_;
using base=mod_ops<dynamic_modnum,uint32_t>;
using base::base;
using base::v;
using base::reduce;
static void set_mod(int mod){
assert(mod>0);
MOD_=uint32_t(mod);
BARRETT_M=(uint64_t(-1)/MOD);
}
static uint32_t barrett_reduce_partial(uint64_t a){
return uint32_t(a-uint64_t((__uint128_t(BARRETT_M)*a)>>64)*MOD);
}
static uint32_t barrett_reduce(uint64_t a){
int32_t res=int32_t(barrett_reduce_partial(a)-MOD);
return uint32_t((res<0)?res+int32_t(MOD):res);
}
struct mod_reader{
friend std::istream&operator>>(std::istream&i,mod_reader){
int mod;i>>mod;
Self::set_mod(mod);
return i;
}
};
static mod_reader MOD_READER(){
return mod_reader();
}
static uint32_t reduce(std::unsigned_integral auto x){
static_assert(sizeof(x)<=8);
return barrett_reduce(x);
}
explicit operator int()const{return int(v);}
Self&operator*=(const Self&o){
v=barrett_reduce(uint64_t(v)*o.v);
return*this;
}
};
template<typename T>struct mod_constraint{
T v,mod;
friend mod_constraint operator&(mod_constraint a,mod_constraint b){
if(a.mod<b.mod)std::swap(a,b);
if(b.mod==1)return a;
extended_gcd_result<T>egcd=extended_gcd<T>(a.mod,b.mod);
assert(a.v%egcd.gcd==b.v%egcd.gcd);
T extra=b.v-a.v%b.mod;
extra/=egcd.gcd;
extra*=egcd.coeff_a;
extra%=b.mod/egcd.gcd;
extra+=(extra<0)?b.mod/egcd.gcd:0;
return mod_constraint{
a.v+extra*a.mod,
a.mod*(b.mod/egcd.gcd)
};
}
};
// src/fft/core.hpp
namespace ecnerwala::fft{
template<typename dbl>struct cplx{
dbl x,y;
cplx(dbl x_=0,dbl y_=0):x(x_),y(y_){}
friend cplx operator+(cplx a,cplx b){return cplx(a.x+b.x,a.y+b.y);}
friend cplx operator-(cplx a,cplx b){return cplx(a.x-b.x,a.y-b.y);}
friend cplx operator*(cplx a,cplx b){return cplx(a.x*b.x-a.y*b.y,a.x*b.y+a.y*b.x);}
friend cplx conj(cplx a){return cplx(a.x,-a.y);}
friend cplx inv(cplx a){dbl n=(a.x*a.x+a.y*a.y);return cplx(a.x/n,-a.y/n);}
};
template<typename num>struct getRoot{
static num f(int k)=delete;
};
template<typename dbl>struct getRoot<cplx<dbl>>{
static cplx<dbl>f(int k){
#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif
dbl a=2*M_PI/k;
return cplx<dbl>(cos(a),sin(a));
}
};
template<int MOD>struct primitive_root{
static const int value;
};
template<>struct primitive_root<998244353>{
static const int value=3;
};
template<>struct primitive_root<(15<<27)+1>{
static const int value=31;
};
template<>struct primitive_root<(127<<24)+1>{
static const int value=3;
};
template<>struct primitive_root<(7<<26)+1>{
static const int value=3;
};
template<>struct primitive_root<(5<<25)+1>{
static const int value=3;
};
template<int MOD>struct getRoot<modnum<MOD>>{
static modnum<MOD>f(int k){
assert((MOD-1)%k==0);
return power(modnum<MOD>(primitive_root<MOD>::value),(MOD-1)/k);
}
};
template<>struct getRoot<mod_goldilocks>{
static mod_goldilocks f(int k){
assert((mod_goldilocks::MOD-1)%k==0);
return power(mod_goldilocks(mod_goldilocks::PRIMITIVE_ROOT),(mod_goldilocks::MOD-1)/k);
}
};
template<typename num>struct fft_core{
static inline vector<int>rev;
static inline vector<num>rt,inv_rt;
static void init(int n){
if(n<=sz(rt))return;
rev.resize(n);
for(int i=0;i<n;i++){
rev[i]=(rev[i>>1]|((i&1)*n))>>1;
}
rt.reserve(n);inv_rt.reserve(n);
while(sz(rt)<2&&sz(rt)<n){rt.push_back(num(1));inv_rt.push_back(num(1));}
for(int k=sz(rt);k<n;k*=2){
rt.resize(2*k);inv_rt.resize(2*k);
num z=getRoot<num>::f(2*k);
num iz=inv(z);
for(int i=k/2;i<k;i++){
rt[2*i]=rt[i],rt[2*i+1]=rt[i]*z;
inv_rt[2*i]=inv_rt[i],inv_rt[2*i+1]=inv_rt[i]*iz;
}
}
}
static int brev(int i,int n){
int s=__builtin_ctz(unsigned(sz(rev)/n));
return rev[i]>>s;
}
static int conj_index(int j){
return j==0?0:j^((1<<(31-__builtin_clz(unsigned(j))))-1);
}
static void forward(std::span<num>a){
int n=sz(a);
if(n<=1)return;
init(n);
for(int k=n/2;k>=1;k/=2){
for(int i=0;i<n;i+=2*k){
for(int j=0;j<k;j++){
num u=a[i+j],v=a[i+j+k];
a[i+j]=u+v;
a[i+j+k]=(u-v)*rt[j+k];
}
}
}
}
static void inverse(std::span<num>a){
int n=sz(a);
if(n<=1)return;
init(n);
for(int k=1;k<n;k*=2){
for(int i=0;i<n;i+=2*k){
for(int j=0;j<k;j++){
num t=inv_rt[j+k]*a[i+j+k];
a[i+j+k]=a[i+j]-t;
a[i+j]=a[i+j]+t;
}
}
}
}
static void extend(std::span<num>t,std::span<const num>coeffs){
int n=sz(t)/2;
assert(sz(coeffs)<=2*n);
init(sz(t));
auto b=t.subspan(n,n);
int lo=min(sz(coeffs),n);
for(int i=0;i<lo;i++){
b[i]=coeffs[i]*rt[n+i];
}
std::fill(b.begin()+lo,b.end(),num(0));
for(int i=n;i<sz(coeffs);i++){
b[i-n]=b[i-n]-coeffs[i]*rt[i];
}
forward(b);
}
static void even_half(std::span<const num>t,std::span<num>out){
int n=sz(out);
assert(sz(t)>=2*n);
num half=inv(num(2));
for(int j=0;j<n;j++)out[j]=(t[2*j]+t[2*j+1])*half;
}
static void odd_half(std::span<const num>t,std::span<num>out){
int n=sz(out);
assert(sz(t)>=2*n);
init(2*n);
num half=inv(num(2));
for(int j=0;j<n;j++){
out[j]=(t[2*j]-t[2*j+1])*half*inv_rt[n+brev(j,n)];
}
}
};
}
// src/fft/engines/ntt.hpp
namespace ecnerwala::fft::engines{
template<typename num>struct ntt{
using value_type=num;
static constexpr bool commutative=true;
using core=fft_core<num>;
struct transformed{
vector<num>v;
int size()const{return sz(v);}
};
using product=transformed;
static constexpr int unit_scale=0;
template<int A=0>using transformed_t=transformed;
template<int K=0>using product_t=product;
static transformed transform(std::span<const num>a,int n){
assert(sz(a)<=2*n);
transformed r;
r.v.assign(n,num(0));
int lo=min(sz(a),n);
std::copy(a.begin(),a.begin()+lo,r.v.begin());
for(int i=n;i<sz(a);i++)r.v[i-n]+=a[i];
core::forward(std::span<num>(r.v));
return r;
}
static void extend_to(transformed&t,int m,std::span<const num>coeffs){
assert(!(m&(m-1))&&sz(coeffs)<=2*m);
if(t.size()>=m)return;
if(t.size()==0){t=transform(coeffs,m);return;}
while(t.size()<m){
int s=t.size();
t.v.resize(2*s);
core::extend(std::span<num>(t.v),coeffs.first(size_t(min(sz(coeffs),2*s))));
}
}
static transformed downsample(const transformed&t,int n,bool odd){
transformed r;r.v.resize(n);
if(odd)core::odd_half(std::span<const num>(t.v),std::span<num>(r.v));
else core::even_half(std::span<const num>(t.v),std::span<num>(r.v));
return r;
}
static transformed negate_arg(const transformed&t,int n){
assert(n>=2&&t.size()>=n);
transformed r;r.v.resize(n);
for(int j=0;j<n;j++)r.v[j]=t.v[j^1];
return r;
}
static product mul(const transformed&a,const transformed&b,int n){
assert(a.size()>=n&&b.size()>=n);
product p;p.v.resize(n);
for(int i=0;i<n;i++)p.v[i]=a.v[i]*b.v[i];
return p;
}
static product sq(const transformed&a,int n){return mul(a,a,n);}
static product mul2(
const transformed&a1,const transformed&b1,
const transformed&a2,const transformed&b2,
int n
){
assert(a1.size()>=n&&b1.size()>=n&&a2.size()>=n&&b2.size()>=n);
product p;p.v.resize(n);
for(int i=0;i<n;i++)p.v[i]=a1.v[i]*b1.v[i]+a2.v[i]*b2.v[i];
return p;
}
static product add(product&&a,const product&b){
assert(a.size()==b.size());
for(int i=0;i<a.size();i++)a.v[i]+=b.v[i];
return std::move(a);
}
template<typename Op=assign_op>static void finish(product&&p,std::span<num>out,Op op={}){
int n=p.size();
assert(sz(out)<=n);
core::inverse(std::span<num>(p.v));
num d=inv(num(n));
for(int i=0;i<sz(out);i++)op(out[i],p.v[i]*d);
}
};
}
// src/fft/engines/real.hpp
namespace ecnerwala::fft::engines{
template<typename dbl=double>struct real{
using value_type=dbl;
static constexpr bool commutative=true;
using cnum=cplx<dbl>;
using core=fft_core<cnum>;
struct transformed{
vector<cnum>v;
int size()const{return 2*sz(v);}
};
using product=transformed;
static constexpr int unit_scale=0;
template<int A=0>using transformed_t=transformed;
template<int K=0>using product_t=product;
static int packed_size(int n){return std::max(n/2,1);}
static void pack(std::span<const dbl>a,std::span<cnum>c){
for(int i=0;i<sz(a);i++)(i&1?c[i/2].y:c[i/2].x)=a[i];
}
static cnum part(const transformed&f,int t,bool odd){
cnum g=conj(f.v[core::conj_index(t)]);
return odd?(f.v[t]-g)*cnum(0,dbl(-0.5)):(f.v[t]+g)*cnum(dbl(0.5));
}
static cnum retangle(cnum s0,cnum s1,int mo,int q){
cnum s=(s0+s1)*cnum(dbl(0.5));
cnum d=(s0-s1)*cnum(dbl(0.5))*core::inv_rt[mo+q];
return s+cnum(-d.y,d.x);
}
static transformed transform(std::span<const dbl>a,int n){
assert(sz(a)<=2*n);
transformed r;
r.v.assign(packed_size(n),cnum(0));
for(int i=0;i<sz(a);i++){
int j=i<n?i:i-n;
((j&1)?r.v[j/2].y:r.v[j/2].x)+=a[i];
}
core::forward(std::span<cnum>(r.v));
return r;
}
static void extend_to(transformed&t,int m,std::span<const dbl>coeffs){
assert(!(m&(m-1))&&sz(coeffs)<=2*m);
if(t.size()>=m)return;
if(t.size()==0){t=transform(coeffs,m);return;}
auto buf=buffer_pool<cnum>::get((sz(coeffs)+1)/2);
std::fill(buf.span().begin(),buf.span().end(),cnum(0));
pack(coeffs,buf.span());
while(t.size()<m){
int s=sz(t.v);
t.v.resize(2*s);
core::extend(
std::span<cnum>(t.v),
std::span<const cnum>(buf.span()).first(size_t(min(sz(buf.span()),2*s)))
);
}
}
static transformed downsample(const transformed&t,int n,bool odd){return half(t,n,odd);}
static transformed negate_arg(const transformed&t,int n){
int m=packed_size(n);
assert(n>=2&&sz(t.v)>=m);
transformed r;r.v.resize(m);
for(int j=0;j<m;j++)r.v[j]=conj(t.v[core::conj_index(j)]);
return r;
}
static transformed half(const transformed&f,int n,bool odd){
assert(n>=2&&f.size()>=2*n);
int mo=n/2;
core::init(2*mo);
transformed r;r.v.resize(mo);
for(int u=0;u<mo;u++){
r.v[u]=retangle(part(f,2*u,odd),part(f,2*u+1,odd),mo,core::brev(u,mo));
}
return r;
}
static product mul(const transformed&a,const transformed&b,int n){
int m=packed_size(n);
assert(a.size()>=n&&b.size()>=n);
core::init(2*m);
product p;p.v.resize(m);
for(int t=0;t<m;t++){
int k=core::brev(t,m);
cnum w=core::rt[m+k];
cnum xa=part(a,t,false),ya=part(a,t,true);
cnum xb=part(b,t,false),yb=part(b,t,true);
cnum p0=(xa+w*ya)*(xb+w*yb);
cnum p1=(xa-w*ya)*(xb-w*yb);
p.v[t]=retangle(p0,p1,m,k);
}
return p;
}
static product sq(const transformed&a,int n){return mul(a,a,n);}
static product mul2(
const transformed&a1,const transformed&b1,
const transformed&a2,const transformed&b2,
int n
){
int m=packed_size(n);
assert(a1.size()>=n&&b1.size()>=n&&a2.size()>=n&&b2.size()>=n);
core::init(2*m);
product p;p.v.resize(m);
for(int t=0;t<m;t++){
int k=core::brev(t,m);
cnum w=core::rt[m+k];
cnum xa1=part(a1,t,false),ya1=part(a1,t,true);
cnum xb1=part(b1,t,false),yb1=part(b1,t,true);
cnum xa2=part(a2,t,false),ya2=part(a2,t,true);
cnum xb2=part(b2,t,false),yb2=part(b2,t,true);
cnum p0=(xa1+w*ya1)*(xb1+w*yb1)+(xa2+w*ya2)*(xb2+w*yb2);
cnum p1=(xa1-w*ya1)*(xb1-w*yb1)+(xa2-w*ya2)*(xb2-w*yb2);
p.v[t]=retangle(p0,p1,m,k);
}
return p;
}
static product add(product&&a,const product&b){
assert(a.size()==b.size());
for(int i=0;i<sz(a.v);i++)a.v[i]=a.v[i]+b.v[i];
return std::move(a);
}
template<typename Op=assign_op>static void finish(product&&p,std::span<dbl>out,Op op={}){
int m=sz(p.v);
assert(sz(out)<=2*m);
core::inverse(std::span<cnum>(p.v));
dbl d=dbl(1)/dbl(m);
for(int i=0;i<sz(out);i++)op(out[i],(i&1?p.v[i/2].y:p.v[i/2].x)*d);
}
};
}
// src/fft/engines/split.hpp
namespace ecnerwala::fft::engines{
template<typename mnum>struct split{
static_assert(sizeof(decltype(mnum::MOD))<=4,"limbs must fit 15 bits");
using value_type=mnum;
static constexpr bool commutative=true;
static constexpr int unit_scale=1;
using cnum=cplx<double>;
using core=fft_core<cnum>;
template<int A=1>struct transformed_t{
vector<cnum>v;
int size()const{return sz(v);}
transformed_t()=default;
explicit transformed_t(vector<cnum>&&v_):v(std::move(v_)){}
template<int A2>requires(A2!=A)explicit(A2>A)transformed_t(transformed_t<A2>&&o)
:v(std::move(o.v)){}
};
using transformed=transformed_t<1>;
template<int K>struct product_t{
vector<cnum>lo,hi;
int size()const{return sz(lo);}
product_t()=default;
product_t(vector<cnum>&&lo_,vector<cnum>&&hi_):lo(std::move(lo_)),hi(std::move(hi_)){}
template<int K2>requires(K2!=K)explicit(K2>K)product_t(product_t<K2>&&o)
:lo(std::move(o.lo)),hi(std::move(o.hi)){}
};
using product=product_t<1>;
static cnum pack(mnum x){
int64_t v=x.balanced();
int64_t hi=(v+(1<<14))>>15;
return cnum(double(v-(hi<<15)),double(hi));
}
static transformed transform(std::span<const mnum>a,int n){
assert(sz(a)<=2*n);
transformed r;
r.v.assign(n,cnum(0));
for(int i=0;i<sz(a);i++){
int j=i<n?i:i-n;
r.v[j]=r.v[j]+pack(a[i]);
}
core::forward(std::span<cnum>(r.v));
return r;
}
static void extend_to(transformed&t,int m,std::span<const mnum>coeffs){
assert(!(m&(m-1))&&sz(coeffs)<=2*m);
if(t.size()>=m)return;
if(t.size()==0){t=transform(coeffs,m);return;}
auto buf=buffer_pool<cnum>::get(sz(coeffs));
for(int i=0;i<sz(coeffs);i++)buf[i]=pack(coeffs[i]);
while(t.size()<m){
int s=t.size();
t.v.resize(2*s);
core::extend(
std::span<cnum>(t.v),
std::span<const cnum>(buf.span()).first(size_t(min(sz(coeffs),2*s)))
);
}
}
static void downsample_core(std::span<const cnum>in,std::span<cnum>out,bool odd){
if(odd)core::odd_half(in,out);
else core::even_half(in,out);
}
template<int A>static transformed_t<A>downsample(const transformed_t<A>&t,int n,bool odd){
transformed_t<A>r;r.v.resize(n);
downsample_core(std::span<const cnum>(t.v),std::span<cnum>(r.v),odd);
return r;
}
template<int K>static product_t<K>downsample(const product_t<K>&p,int n,bool odd){
product_t<K>r;r.lo.resize(n);r.hi.resize(n);
downsample_core(std::span<const cnum>(p.lo),std::span<cnum>(r.lo),odd);
downsample_core(std::span<const cnum>(p.hi),std::span<cnum>(r.hi),odd);
return r;
}
template<int A>static transformed_t<A>negate_arg(const transformed_t<A>&t,int n){
assert(n>=2&&t.size()>=n);
transformed_t<A>r;r.v.resize(n);
for(int j=0;j<n;j++)r.v[j]=t.v[j^1];
return r;
}
template<int A,int B>static transformed_t<A+B>add(transformed_t<A>&&a,const transformed_t<B>&b){
transformed_t<A+B>r{std::move(a.v)};
add_into(r.v,b.v);
return r;
}
static void mul_impl(const vector<cnum>&a,const vector<cnum>&b,vector<cnum>&lo,vector<cnum>&hi,int n,bool acc=false){
core::init(n);
lo.resize(n);hi.resize(n);
for(int i=0;i<n;i++){
int ci=core::conj_index(i);
cnum g0=(b[i]+conj(b[ci]))*cnum(0.5);
cnum t=(b[i]-conj(b[ci]))*cnum(0.5);
cnum g1=cnum(t.y,-t.x);
if(acc){
lo[i]=lo[i]+a[i]*g0;
hi[i]=hi[i]+a[i]*g1;
}else{
lo[i]=a[i]*g0;
hi[i]=a[i]*g1;
}
}
}
template<int A,int B>static product_t<A*B>mul(const transformed_t<A>&a,const transformed_t<B>&b,int n){
assert(a.size()>=n&&b.size()>=n);
product_t<A*B>p;
mul_impl(a.v,b.v,p.lo,p.hi,n);
return p;
}
template<int A>static product_t<A*A>sq(const transformed_t<A>&a,int n){return mul(a,a,n);}
template<int A1,int B1,int A2,int B2>
static product_t<A1*B1+A2*B2>mul2(
const transformed_t<A1>&a1,const transformed_t<B1>&b1,
const transformed_t<A2>&a2,const transformed_t<B2>&b2,
int n
){
assert(a1.size()>=n&&b1.size()>=n&&a2.size()>=n&&b2.size()>=n);
product_t<A1*B1+A2*B2>p;
mul_impl(a1.v,b1.v,p.lo,p.hi,n);
mul_impl(a2.v,b2.v,p.lo,p.hi,n,true);
return p;
}
static void add_into(vector<cnum>&a,const vector<cnum>&b){
assert(sz(a)==sz(b));
for(int i=0;i<sz(a);i++)a[i]=a[i]+b[i];
}
template<int K1,int K2>static product_t<K1+K2>add(product_t<K1>&&a,product_t<K2>&&b){
product_t<K1+K2>r{std::move(a.lo),std::move(a.hi)};
add_into(r.lo,b.lo);
add_into(r.hi,b.hi);
return r;
}
template<int K=1,typename Op=assign_op>static void finish(product_t<K>&&p,std::span<mnum>out,Op op={}){
static_assert(K<=2,"split: accumulated scale too large");
int n=p.size();
assert(sz(out)<=n);
core::inverse(std::span<cnum>(p.lo));
core::inverse(std::span<cnum>(p.hi));
const int64_t m=mnum::MOD;
double d=1.0/double(n);
for(int i=0;i<sz(out);i++){
int64_t v=(llround(p.lo[i].x*d)
+(llround(p.lo[i].y*d)%m<<15)
+(llround(p.hi[i].x*d)%m<<15)
+(llround(p.hi[i].y*d)%m<<30))%m;
if(v<0)v+=m;
op(out[i],mnum(v));
}
}
};
}
// src/fft/engines/crt.hpp
namespace ecnerwala::fft::engines{
template<typename mnum,typename num1=mod_goldilocks,typename num2=modnum<(15<<27)+1>>
struct crt{
static_assert(sizeof(decltype(mnum::MOD))<=4,"n (MOD/2)^2 must fit the CRT modulus product");
using value_type=mnum;
static constexpr bool commutative=true;
static constexpr int unit_scale=1;
using E1=ntt<num1>;
using E2=ntt<num2>;
template<int A=1>struct transformed_t{
typename E1::transformed t1;
typename E2::transformed t2;
int size()const{return t1.size();}
transformed_t()=default;
transformed_t(typename E1::transformed&&t1_,typename E2::transformed&&t2_)
:t1(std::move(t1_)),t2(std::move(t2_)){}
template<int A2>requires(A2!=A)explicit(A2>A)transformed_t(transformed_t<A2>&&o)
:t1(std::move(o.t1)),t2(std::move(o.t2)){}
};
using transformed=transformed_t<1>;
template<int K>struct product_t{
typename E1::product p1;
typename E2::product p2;
int size()const{return sz(p1);}
product_t()=default;
product_t(typename E1::product&&p1_,typename E2::product&&p2_)
:p1(std::move(p1_)),p2(std::move(p2_)){}
template<int K2>requires(K2!=K)explicit(K2>K)product_t(product_t<K2>&&o)
:p1(std::move(o.p1)),p2(std::move(o.p2)){}
};
using product=product_t<1>;
static transformed transform(std::span<const mnum>a,int n){
assert(sz(a)<=2*n);
auto b1=buffer_pool<num1>::get(sz(a));
auto b2=buffer_pool<num2>::get(sz(a));
for(int i=0;i<sz(a);i++){int64_t v=a[i].balanced();b1[i]=num1(v);b2[i]=num2(v);}
return transformed{
E1::transform(std::span<const num1>(b1.span()),n),
E2::transform(std::span<const num2>(b2.span()),n),
};
}
static void extend_to(transformed&t,int m,std::span<const mnum>coeffs){
if(t.size()>=m)return;
auto b1=buffer_pool<num1>::get(sz(coeffs));
auto b2=buffer_pool<num2>::get(sz(coeffs));
for(int i=0;i<sz(coeffs);i++){int64_t v=coeffs[i].balanced();b1[i]=num1(v);b2[i]=num2(v);}
E1::extend_to(t.t1,m,std::span<const num1>(b1.span()));
E2::extend_to(t.t2,m,std::span<const num2>(b2.span()));
}
template<int A>static transformed_t<A>downsample(const transformed_t<A>&t,int n,bool odd){
return transformed_t<A>{E1::downsample(t.t1,n,odd),E2::downsample(t.t2,n,odd)};
}
template<int K>static product_t<K>downsample(const product_t<K>&p,int n,bool odd){
return product_t<K>{E1::downsample(p.p1,n,odd),E2::downsample(p.p2,n,odd)};
}
template<int A>static transformed_t<A>negate_arg(const transformed_t<A>&t,int n){
return transformed_t<A>{E1::negate_arg(t.t1,n),E2::negate_arg(t.t2,n)};
}
template<int A,int B>static transformed_t<A+B>add(transformed_t<A>&&a,const transformed_t<B>&b){
return transformed_t<A+B>{E1::add(std::move(a.t1),b.t1),E2::add(std::move(a.t2),b.t2)};
}
template<int A,int B>static product_t<A*B>mul(const transformed_t<A>&a,const transformed_t<B>&b,int n){
return product_t<A*B>{E1::mul(a.t1,b.t1,n),E2::mul(a.t2,b.t2,n)};
}
template<int A>static product_t<A*A>sq(const transformed_t<A>&a,int n){return mul(a,a,n);}
template<int A1,int B1,int A2,int B2>
static product_t<A1*B1+A2*B2>mul2(
const transformed_t<A1>&a1,const transformed_t<B1>&b1,
const transformed_t<A2>&a2,const transformed_t<B2>&b2,
int n
){
return product_t<A1*B1+A2*B2>{
E1::mul2(a1.t1,b1.t1,a2.t1,b2.t1,n),
E2::mul2(a1.t2,b1.t2,a2.t2,b2.t2,n),
};
}
template<int K1,int K2>static product_t<K1+K2>add(product_t<K1>&&a,product_t<K2>&&b){
return product_t<K1+K2>{E1::add(std::move(a.p1),b.p1),E2::add(std::move(a.p2),b.p2)};
}
template<int K=1,typename Op=assign_op>static void finish(product_t<K>&&p,std::span<mnum>out,Op op={}){
static_assert(K<=2,"crt: accumulated scale too large");
int n=p.size();
assert(sz(out)<=n);
auto o1=buffer_pool<num1>::get(sz(out));
auto o2=buffer_pool<num2>::get(sz(out));
E1::finish(std::move(p.p1),o1.span());
E2::finish(std::move(p.p2),o2.span());
num1 inv_n2=inv(num1(num2::MOD));
num2 inv_n1=inv(num2(num1::MOD));
__int128_t whole=__int128_t(num1::MOD)*__int128_t(num2::MOD);
mnum m1_mod=mnum(num1::MOD);
mnum m2_mod=mnum(num2::MOD);
mnum whole_mod=m1_mod*m2_mod;
for(int i=0;i<sz(out);i++){
num1 v1=o1[i]*inv_n2;
num2 v2=o2[i]*inv_n1;
mnum o_mod=mnum(uint64_t(v1))*m2_mod+mnum(int(v2))*m1_mod;
__int128_t o_exact=__int128_t(uint64_t(v1))*__int128_t(num2::MOD)+__int128_t(int(v2))*__int128_t(num1::MOD);
if(o_exact>=whole){o_exact-=whole;o_mod-=whole_mod;}
if(o_exact>whole/2)o_mod-=whole_mod;
op(out[i],o_mod);
}
}
};
}
// src/fft/test_util.test.hpp
namespace ecnerwala{
namespace fft{
template<typename T>std::vector<T>multiply_slow(const std::vector<T>&a,const std::vector<T>&b){
if(a.empty()||b.empty())return{};
std::vector<T>res(a.size()+b.size()-1);
for(int i=0;i<int(a.size());i++){
for(int j=0;j<int(b.size());j++){
res[i+j]+=a[i]*b[j];
}
}
return res;
}
template<typename T>T rnd_val(std::mt19937&mt){
if constexpr(std::is_floating_point_v<T>)return T(int(mt()%1024));
else return T(mt());
}
template<typename T>void fill_rnd(std::vector<T>&v,std::mt19937&mt){
for(T&x:v)x=rnd_val<T>(mt);
}
template<typename T>void check_eq(const std::vector<T>&got,const std::vector<T>&want){
REQUIRE(got.size()==want.size());
for(int i=0;i<int(got.size());i++){
INFO("i = "<<i);
if constexpr(std::is_floating_point_v<T>)REQUIRE(llround(got[i])==llround(want[i]));
else REQUIRE(got[i]==want[i]);
}
}
#define ALL_ENGINES \
engines::ntt<modnum<998244353>>,engines::ntt<mod_goldilocks>,engines::real<double>,\
engines::split<modnum<int(1e9)+7>>,engines::crt<modnum<int(1e9)+7>>
#define MOD_ENGINES \
engines::ntt<modnum<998244353>>,engines::ntt<mod_goldilocks>,\
engines::split<modnum<int(1e9)+7>>,engines::crt<modnum<int(1e9)+7>>
}}
// src/fft/poly.test.cpp
namespace ecnerwala{
namespace fft{
using namespace std;
TEST_CASE("poly::form evaluation and transposed multiplication","[fft]"){
using num=modnum<998244353>;
using E=engines::ntt<num>;
mt19937 mt(Catch::getSeed());
int n=40;
num z=num(mt());
auto f=poly::form<E>::polynomial_evaluation(z,n);
vector<num>sv(30);
fill_rnd(sv,mt);
poly::vec<E>s((span<const num>(sv)));
REQUIRE(f(s)==s(z));
vector<num>qv(11),s2v(n-11);
fill_rnd(qv,mt);
fill_rnd(s2v,mt);
poly::vec<E>q((span<const num>(qv))),s2((span<const num>(s2v)));
auto fq=f.composed_with(q);
REQUIRE(fq(s2)==f(s2*q));
REQUIRE(fq(s2)==(s2*q)(z));
series::trunc<E>t(size_t(n),num{});
fill_rnd(t,mt);
auto ft=f.composed_with(t);
REQUIRE(ft.len()==n);
for(int j=0;j<n;j++){
num want{};
for(int d=0;d<=j;d++)want+=t[d]*f.rev_series()[j-d];
REQUIRE(ft.rev_series()[j]==want);
}
series::exact<E>e(t.begin(),t.begin()+11);
auto fe=f.composed_with(e);
REQUIRE(fe.len()==n);
for(int j=0;j<n;j++){
num want{};
for(int d=0;d<=j&&d<11;d++)want+=t[d]*f.rev_series()[j-d];
REQUIRE(fe.rev_series()[j]==want);
}
}
TEST_CASE("poly::vec reversed storage and series interop","[fft]"){
using num=modnum<998244353>;
using E=engines::ntt<num>;
mt19937 mt(Catch::getSeed());
vector<num>pa(37),pb(23);
fill_rnd(pa,mt);
fill_rnd(pb,mt);
poly::vec<E>a((span<const num>(pa))),b((span<const num>(pb)));
REQUIRE(a[0]==pa[0]);
REQUIRE(a.leading()==pa[36]);
REQUIRE(a.rev_series()[0]==pa[36]);
REQUIRE(a.rev_series()[36]==pa[0]);
poly::vec<E>p=a*b;
check_eq(vector<num>(p.begin(),p.end()),multiply_slow(pa,pb));
REQUIRE(square(a)==a*a);
num x=num(mt());
REQUIRE(p(x)==a(x)*b(x));
poly::vec<E>s=a+b,d=b-a;
REQUIRE(s.len()==37);
for(int i=0;i<37;i++)REQUIRE(s[i]==pa[i]+(i<23?pb[i]:num(0)));
for(int i=0;i<37;i++)REQUIRE(d[i]==(i<23?pb[i]:num(0))-pa[i]);
poly::vec<E>g=a;
g.shift(2);
REQUIRE(g.len()==39);
REQUIRE(g[0]==num(0));
REQUIRE(g[1]==num(0));
for(int i=0;i<37;i++)REQUIRE(g[i+2]==pa[i]);
REQUIRE(g.rev_series().data()[0]==pa[36]);
const series::exact<E>&ra=a.rev_series();
REQUIRE(ra.len()==37);
for(int i=0;i<37;i++)REQUIRE(ra[i]==pa[36-i]);
REQUIRE(poly::vec<E>::from_rev_series(ra)==a);
series::exact<E>xa(a.begin(),a.end());
REQUIRE(equal(xa.begin(),xa.end(),pa.begin(),pa.end()));
REQUIRE(a.unrev_series(10)==series::trunc<E>(pa.begin(),pa.begin()+10));
std::vector<num>vals(60);
fill_rnd(vals,mt);
series::cached_exact<E>cv(series::exact<E>(vals.begin(),vals.end()));
series::cached_exact<E>ca(a.rev_series());
auto mp=middle_product(cv,ca);
auto naive=[&](int j){
num r{};
for(int t=0;t<37;t++)r+=pa[t]*vals[j+t];
return r;
};
for(int j=0;j<sz(mp);j++)REQUIRE(mp[size_t(j)]==naive(j));
}
TEST_CASE("poly::cached products","[fft]"){
using num=modnum<998244353>;
using E=engines::ntt<num>;
mt19937 mt(Catch::getSeed());
vector<num>pa(37),pb(23);
fill_rnd(pa,mt);
fill_rnd(pb,mt);
poly::vec<E>a((span<const num>(pa))),b((span<const num>(pb)));
auto p=a*b;
static_assert(std::is_same_v<decltype(p),poly::cached<E>>);
REQUIRE(p.rev_series().cache().size()>0);
poly::vec<E>pp=a*b;
check_eq(vector<num>(pp.begin(),pp.end()),multiply_slow(pa,pb));
poly::cached<E>ca(a),cb(b);
REQUIRE(ca==a);
REQUIRE(ca*cb==p);
REQUIRE(ca*b==p);
auto sq=square(ca);
static_assert(std::is_same_v<decltype(sq),poly::cached<E>>);
REQUIRE(sq==a*a);
num x=num(mt());
REQUIRE(p(x)==ca(x)*cb(x));
poly::vec<E>q=std::move(ca);
REQUIRE(q==a);
}
TEST_CASE("poly::multipoint and poly::interpolate","[fft]"){
using num=modnum<998244353>;
mt19937 mt(Catch::getSeed());
for(int n:{1,2,3,8,17,40}){
INFO("n = "<<n);
vector<num>coeffs(n);
for(num&x:coeffs){x=num(mt());}
poly::vec<engines::ntt<num>>p((span<const num>(coeffs)));
vector<num>pts(n);
for(int i=0;i<n;i++)pts[i]=num(1000+i);
auto vals=poly::multipoint<engines::ntt<num>>(p,pts);
for(int i=0;i<n;i++){
REQUIRE(vals[i]==p(pts[i]));
}
auto rec=poly::interpolate<engines::ntt<num>>(pts,vals);
REQUIRE(rec==p);
}
}
}}
#pragma GCC diagnostic pop
// clang-format on
// @formatter:on