ecnerwala's competitive programming library
#include "fft/ap_sampled_poly.hpp"
| Coverage | Exec / Excl / Total | |
|---|---|---|
| Lines | 68.6% | 72 / 0 / 105 |
| Functions | 100.0% | 2 / 0 / 2 |
| Branches | 81.0% | 119 / 0 / 147 |
| Full report |
#pragma once
#include <cassert>
#include <span>
#include <vector>
#include "fft/multiply.hpp"
namespace ecnerwala {
// A polynomial represented by its values evaluated at an Arithmetic Progression (AP).
// TODO: The AP is always assumed to be 0..length-1; store an explicit offset/gap instead?
// Maybe not, this is just more convenient.
template <fft::engine E>
struct ap_sampled_poly : public std::vector<typename E::value_type> {
using T = typename E::value_type;
using std::vector<T>::vector;
int len() const {
return int(this->size());
}
int degree() const {
return len() - 1;
}
ap_sampled_poly& operator += (const ap_sampled_poly& o) {
assert(len() == o.len());
for (int i = 0; i < int(o.size()); i++) {
(*this)[i] += o[i];
}
return *this;
}
friend ap_sampled_poly operator + (const ap_sampled_poly& a, const ap_sampled_poly& b) {
assert(a.size() == b.size());
ap_sampled_poly r(a.size());
for (int i = 0; i < r.len(); i++) {
r[i] = a[i] + b[i];
}
return r;
}
ap_sampled_poly& operator -= (const ap_sampled_poly& o) {
assert(len() == o.len());
for (int i = 0; i < int(o.size()); i++) {
(*this)[i] -= o[i];
}
return *this;
}
friend ap_sampled_poly operator - (const ap_sampled_poly& a, const ap_sampled_poly& b) {
assert(a.size() == b.size());
ap_sampled_poly r(a.size());
for (int i = 0; i < r.len(); i++) {
r[i] = a[i] - b[i];
}
return r;
}
T eval_at(T k) {
if (0 <= int(k) && int(k) < len()) {
return (*this)[int(k)];
} else {
// Just do the lagrange interpolation
std::vector<T> terms(*this);
{
// Inverse factorial terms
T v = 1;
for (int i = 1; i <= len(); i++) v *= T(i);
v = inv(v);
for (int i = len()-1; i >= 0; i--) {
v *= T(i+1);
terms[i] *= v;
terms[len()-1-i] *= (i & 1) ? -v : v;
}
}
{
// Prefix terms
T v = 1;
for (int i = 0; i < len(); i++) {
terms[i] *= v;
v *= T(k - i);
}
}
{
// Suffix terms
T v = 1;
for (int i = len() - 1; i >= 0; i--) {
terms[i] *= v;
v *= T(k - i);
}
}
T res = 0;
for (int i = 0; i < len(); i++) res += terms[i];
return res;
}
}
ap_sampled_poly eval_range(T k, int osz) {
if (osz == 0) {
return ap_sampled_poly(osz);
}
if (len() == 0) {
return ap_sampled_poly(osz, T(0));
}
// Check for overlaps. We're checking in linear time to avoid unpacking T, but it should be plenty fast.
// If the field is very very small and we wrap around several times, our runtime can be bad...
// but then something has already gone wrong, why are you evaluating so many points???
for (int i = -(len() - 1); i <= osz - 1; i++) {
if (k+i == T(0)) {
// everything from [i,i+len()-1) of the output is a match
ap_sampled_poly res; res.reserve(osz);
int lo = std::max(0, i);
int hi = std::min(i+len(), osz);
{
auto pref = eval_range(k, lo);
res.insert(res.end(), pref.begin(), pref.end());
}
res.insert(res.end(), this->begin() + (lo - i), this->begin() + (hi - i));
{
auto suff = eval_range(k + hi, osz - hi);
res.insert(res.end(), suff.begin(), suff.end());
}
return res;
}
}
std::vector<T> inps(*this);
{
// Inverse factorial terms
T v = 1;
for (int i = 1; i <= len(); i++) v *= T(i);
v = inv(v);
for (int i = len()-1; i >= 0; i--) {
v *= T(i+1);
inps[i] *= v;
inps[len()-1-i] *= (i & 1) ? -v : v;
}
}
std::vector<T> inv_offsets(len() + osz - 1);
ap_sampled_poly results(osz);
{
T v = 1;
for (int i = - (len() - 1); i <= osz - 1; i++) {
inv_offsets[i + (len() - 1)] = v;
v *= k + i;
if (i >= 0) results[i] = v;
}
// Assert there's no overlap
assert(v != T(0));
v = inv(v);
for (int i = osz - 1; i >= -(len() - 1); i--) {
inv_offsets[i + (len() - 1)] *= v;
v *= k + i;
if (i + (len() - 1) <= osz - 1) {
results[i + (len() - 1)] *= v;
}
}
}
std::vector<T> prod = fft::middle_product<E>(inv_offsets, inps);
assert(int(prod.size()) == osz);
for (int i = 0; i < osz; i++) results[i] *= prod[i];
return results;
}
void extend_right() {
this->push_back(eval_at(T(len())));
}
void extend_left() {
this->insert(this->begin(), eval_at(T(-1)));
}
[[nodiscard]] ap_sampled_poly prefix_sum_inclusive() const {
ap_sampled_poly r = *this;
r.extend_right();
for (int i = 1; i < r.len(); i++) r[i] += r[i-1];
return r;
}
};
/* namespace ecnerwala */ }
#include <cassert>
#include <span>
#include <vector>
#include <algorithm>
#include <concepts>
#include <cstddef>
#include <utility>
#include <type_traits>
#include <iterator>
#line 2 "src/fft/ap_sampled_poly.hpp"
#line 6 "src/fft/ap_sampled_poly.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 8 "src/fft/ap_sampled_poly.hpp"
namespace ecnerwala {
// A polynomial represented by its values evaluated at an Arithmetic Progression (AP).
// TODO: The AP is always assumed to be 0..length-1; store an explicit offset/gap instead?
// Maybe not, this is just more convenient.
template <fft::engine E>
struct ap_sampled_poly : public std::vector<typename E::value_type> {
using T = typename E::value_type;
using std::vector<T>::vector;
int len() const {
return int(this->size());
}
int degree() const {
return len() - 1;
}
ap_sampled_poly& operator += (const ap_sampled_poly& o) {
assert(len() == o.len());
for (int i = 0; i < int(o.size()); i++) {
(*this)[i] += o[i];
}
return *this;
}
friend ap_sampled_poly operator + (const ap_sampled_poly& a, const ap_sampled_poly& b) {
assert(a.size() == b.size());
ap_sampled_poly r(a.size());
for (int i = 0; i < r.len(); i++) {
r[i] = a[i] + b[i];
}
return r;
}
ap_sampled_poly& operator -= (const ap_sampled_poly& o) {
assert(len() == o.len());
for (int i = 0; i < int(o.size()); i++) {
(*this)[i] -= o[i];
}
return *this;
}
friend ap_sampled_poly operator - (const ap_sampled_poly& a, const ap_sampled_poly& b) {
assert(a.size() == b.size());
ap_sampled_poly r(a.size());
for (int i = 0; i < r.len(); i++) {
r[i] = a[i] - b[i];
}
return r;
}
T eval_at(T k) {
if (0 <= int(k) && int(k) < len()) {
return (*this)[int(k)];
} else {
// Just do the lagrange interpolation
std::vector<T> terms(*this);
{
// Inverse factorial terms
T v = 1;
for (int i = 1; i <= len(); i++) v *= T(i);
v = inv(v);
for (int i = len()-1; i >= 0; i--) {
v *= T(i+1);
terms[i] *= v;
terms[len()-1-i] *= (i & 1) ? -v : v;
}
}
{
// Prefix terms
T v = 1;
for (int i = 0; i < len(); i++) {
terms[i] *= v;
v *= T(k - i);
}
}
{
// Suffix terms
T v = 1;
for (int i = len() - 1; i >= 0; i--) {
terms[i] *= v;
v *= T(k - i);
}
}
T res = 0;
for (int i = 0; i < len(); i++) res += terms[i];
return res;
}
}
ap_sampled_poly eval_range(T k, int osz) {
if (osz == 0) {
return ap_sampled_poly(osz);
}
if (len() == 0) {
return ap_sampled_poly(osz, T(0));
}
// Check for overlaps. We're checking in linear time to avoid unpacking T, but it should be plenty fast.
// If the field is very very small and we wrap around several times, our runtime can be bad...
// but then something has already gone wrong, why are you evaluating so many points???
for (int i = -(len() - 1); i <= osz - 1; i++) {
if (k+i == T(0)) {
// everything from [i,i+len()-1) of the output is a match
ap_sampled_poly res; res.reserve(osz);
int lo = std::max(0, i);
int hi = std::min(i+len(), osz);
{
auto pref = eval_range(k, lo);
res.insert(res.end(), pref.begin(), pref.end());
}
res.insert(res.end(), this->begin() + (lo - i), this->begin() + (hi - i));
{
auto suff = eval_range(k + hi, osz - hi);
res.insert(res.end(), suff.begin(), suff.end());
}
return res;
}
}
std::vector<T> inps(*this);
{
// Inverse factorial terms
T v = 1;
for (int i = 1; i <= len(); i++) v *= T(i);
v = inv(v);
for (int i = len()-1; i >= 0; i--) {
v *= T(i+1);
inps[i] *= v;
inps[len()-1-i] *= (i & 1) ? -v : v;
}
}
std::vector<T> inv_offsets(len() + osz - 1);
ap_sampled_poly results(osz);
{
T v = 1;
for (int i = - (len() - 1); i <= osz - 1; i++) {
inv_offsets[i + (len() - 1)] = v;
v *= k + i;
if (i >= 0) results[i] = v;
}
// Assert there's no overlap
assert(v != T(0));
v = inv(v);
for (int i = osz - 1; i >= -(len() - 1); i--) {
inv_offsets[i + (len() - 1)] *= v;
v *= k + i;
if (i + (len() - 1) <= osz - 1) {
results[i + (len() - 1)] *= v;
}
}
}
std::vector<T> prod = fft::middle_product<E>(inv_offsets, inps);
assert(int(prod.size()) == osz);
for (int i = 0; i < osz; i++) results[i] *= prod[i];
return results;
}
void extend_right() {
this->push_back(eval_at(T(len())));
}
void extend_left() {
this->insert(this->begin(), eval_at(T(-1)));
}
[[nodiscard]] ap_sampled_poly prefix_sum_inclusive() const {
ap_sampled_poly r = *this;
r.extend_right();
for (int i = 1; i < r.len(); i++) r[i] += r[i-1];
return r;
}
};
/* namespace ecnerwala */ }
// 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 <cassert>
// 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/ap_sampled_poly.hpp
namespace ecnerwala{
template<fft::engine E>
struct ap_sampled_poly:public std::vector<typename E::value_type>{
using T=typename E::value_type;
using std::vector<T>::vector;
int len()const{
return int(this->size());
}
int degree()const{
return len()-1;
}
ap_sampled_poly&operator+=(const ap_sampled_poly&o){
assert(len()==o.len());
for(int i=0;i<int(o.size());i++){
(*this)[i]+=o[i];
}
return*this;
}
friend ap_sampled_poly operator+(const ap_sampled_poly&a,const ap_sampled_poly&b){
assert(a.size()==b.size());
ap_sampled_poly r(a.size());
for(int i=0;i<r.len();i++){
r[i]=a[i]+b[i];
}
return r;
}
ap_sampled_poly&operator-=(const ap_sampled_poly&o){
assert(len()==o.len());
for(int i=0;i<int(o.size());i++){
(*this)[i]-=o[i];
}
return*this;
}
friend ap_sampled_poly operator-(const ap_sampled_poly&a,const ap_sampled_poly&b){
assert(a.size()==b.size());
ap_sampled_poly r(a.size());
for(int i=0;i<r.len();i++){
r[i]=a[i]-b[i];
}
return r;
}
T eval_at(T k){
if(0<=int(k)&&int(k)<len()){
return(*this)[int(k)];
}else{
std::vector<T>terms(*this);
{
T v=1;
for(int i=1;i<=len();i++)v*=T(i);
v=inv(v);
for(int i=len()-1;i>=0;i--){
v*=T(i+1);
terms[i]*=v;
terms[len()-1-i]*=(i&1)?-v:v;
}
}
{
T v=1;
for(int i=0;i<len();i++){
terms[i]*=v;
v*=T(k-i);
}
}
{
T v=1;
for(int i=len()-1;i>=0;i--){
terms[i]*=v;
v*=T(k-i);
}
}
T res=0;
for(int i=0;i<len();i++)res+=terms[i];
return res;
}
}
ap_sampled_poly eval_range(T k,int osz){
if(osz==0){
return ap_sampled_poly(osz);
}
if(len()==0){
return ap_sampled_poly(osz,T(0));
}
for(int i=-(len()-1);i<=osz-1;i++){
if(k+i==T(0)){
ap_sampled_poly res;res.reserve(osz);
int lo=std::max(0,i);
int hi=std::min(i+len(),osz);
{
auto pref=eval_range(k,lo);
res.insert(res.end(),pref.begin(),pref.end());
}
res.insert(res.end(),this->begin()+(lo-i),this->begin()+(hi-i));
{
auto suff=eval_range(k+hi,osz-hi);
res.insert(res.end(),suff.begin(),suff.end());
}
return res;
}
}
std::vector<T>inps(*this);
{
T v=1;
for(int i=1;i<=len();i++)v*=T(i);
v=inv(v);
for(int i=len()-1;i>=0;i--){
v*=T(i+1);
inps[i]*=v;
inps[len()-1-i]*=(i&1)?-v:v;
}
}
std::vector<T>inv_offsets(len()+osz-1);
ap_sampled_poly results(osz);
{
T v=1;
for(int i=-(len()-1);i<=osz-1;i++){
inv_offsets[i+(len()-1)]=v;
v*=k+i;
if(i>=0)results[i]=v;
}
assert(v!=T(0));
v=inv(v);
for(int i=osz-1;i>=-(len()-1);i--){
inv_offsets[i+(len()-1)]*=v;
v*=k+i;
if(i+(len()-1)<=osz-1){
results[i+(len()-1)]*=v;
}
}
}
std::vector<T>prod=fft::middle_product<E>(inv_offsets,inps);
assert(int(prod.size())==osz);
for(int i=0;i<osz;i++)results[i]*=prod[i];
return results;
}
void extend_right(){
this->push_back(eval_at(T(len())));
}
void extend_left(){
this->insert(this->begin(),eval_at(T(-1)));
}
[[nodiscard]]ap_sampled_poly prefix_sum_inclusive()const{
ap_sampled_poly r=*this;
r.extend_right();
for(int i=1;i<r.len();i++)r[i]+=r[i-1];
return r;
}
};
}
#pragma GCC diagnostic pop
// clang-format on
// @formatter:on
#pragma once
#include <cassert>
#include <span>
#include <vector>
#include "fft/multiply.hpp"
namespace ecnerwala {
// A polynomial represented by its values evaluated at an Arithmetic Progression (AP).
// TODO: The AP is always assumed to be 0..length-1; store an explicit offset/gap instead?
// Maybe not, this is just more convenient.
template <fft::engine E>
struct ap_sampled_poly : public std::vector<typename E::value_type> {
using T = typename E::value_type;
using std::vector<T>::vector;
int len() const {
return int(this->size());
}
int degree() const {
return len() - 1;
}
ap_sampled_poly& operator += (const ap_sampled_poly& o) {
assert(len() == o.len());
for (int i = 0; i < int(o.size()); i++) {
(*this)[i] += o[i];
}
return *this;
}
friend ap_sampled_poly operator + (const ap_sampled_poly& a, const ap_sampled_poly& b) {
assert(a.size() == b.size());
ap_sampled_poly r(a.size());
for (int i = 0; i < r.len(); i++) {
r[i] = a[i] + b[i];
}
return r;
}
ap_sampled_poly& operator -= (const ap_sampled_poly& o) {
assert(len() == o.len());
for (int i = 0; i < int(o.size()); i++) {
(*this)[i] -= o[i];
}
return *this;
}
friend ap_sampled_poly operator - (const ap_sampled_poly& a, const ap_sampled_poly& b) {
assert(a.size() == b.size());
ap_sampled_poly r(a.size());
for (int i = 0; i < r.len(); i++) {
r[i] = a[i] - b[i];
}
return r;
}
T eval_at(T k) {
if (0 <= int(k) && int(k) < len()) {
return (*this)[int(k)];
} else {
// Just do the lagrange interpolation
std::vector<T> terms(*this);
{
// Inverse factorial terms
T v = 1;
for (int i = 1; i <= len(); i++) v *= T(i);
v = inv(v);
for (int i = len()-1; i >= 0; i--) {
v *= T(i+1);
terms[i] *= v;
terms[len()-1-i] *= (i & 1) ? -v : v;
}
}
{
// Prefix terms
T v = 1;
for (int i = 0; i < len(); i++) {
terms[i] *= v;
v *= T(k - i);
}
}
{
// Suffix terms
T v = 1;
for (int i = len() - 1; i >= 0; i--) {
terms[i] *= v;
v *= T(k - i);
}
}
T res = 0;
for (int i = 0; i < len(); i++) res += terms[i];
return res;
}
}
ap_sampled_poly eval_range(T k, int osz) {
if (osz == 0) {
return ap_sampled_poly(osz);
}
if (len() == 0) {
return ap_sampled_poly(osz, T(0));
}
// Check for overlaps. We're checking in linear time to avoid unpacking T, but it should be plenty fast.
// If the field is very very small and we wrap around several times, our runtime can be bad...
// but then something has already gone wrong, why are you evaluating so many points???
for (int i = -(len() - 1); i <= osz - 1; i++) {
if (k+i == T(0)) {
// everything from [i,i+len()-1) of the output is a match
ap_sampled_poly res; res.reserve(osz);
int lo = std::max(0, i);
int hi = std::min(i+len(), osz);
{
auto pref = eval_range(k, lo);
res.insert(res.end(), pref.begin(), pref.end());
}
res.insert(res.end(), this->begin() + (lo - i), this->begin() + (hi - i));
{
auto suff = eval_range(k + hi, osz - hi);
res.insert(res.end(), suff.begin(), suff.end());
}
return res;
}
}
std::vector<T> inps(*this);
{
// Inverse factorial terms
T v = 1;
for (int i = 1; i <= len(); i++) v *= T(i);
v = inv(v);
for (int i = len()-1; i >= 0; i--) {
v *= T(i+1);
inps[i] *= v;
inps[len()-1-i] *= (i & 1) ? -v : v;
}
}
std::vector<T> inv_offsets(len() + osz - 1);
ap_sampled_poly results(osz);
{
T v = 1;
for (int i = - (len() - 1); i <= osz - 1; i++) {
inv_offsets[i + (len() - 1)] = v;
v *= k + i;
if (i >= 0) results[i] = v;
}
// Assert there's no overlap
assert(v != T(0));
v = inv(v);
for (int i = osz - 1; i >= -(len() - 1); i--) {
inv_offsets[i + (len() - 1)] *= v;
v *= k + i;
if (i + (len() - 1) <= osz - 1) {
results[i + (len() - 1)] *= v;
}
}
}
std::vector<T> prod = fft::middle_product<E>(inv_offsets, inps);
assert(int(prod.size()) == osz);
for (int i = 0; i < osz; i++) results[i] *= prod[i];
return results;
}
void extend_right() {
this->push_back(eval_at(T(len())));
}
void extend_left() {
this->insert(this->begin(), eval_at(T(-1)));
}
[[nodiscard]] ap_sampled_poly prefix_sum_inclusive() const {
ap_sampled_poly r = *this;
r.extend_right();
for (int i = 1; i < r.len(); i++) r[i] += r[i-1];
return r;
}
};
/* namespace ecnerwala */ }