Skip to content

Instantly share code, notes, and snippets.

@lychees

lychees/poly_gen.cpp

Created May 8, 2023 06:21
Show Gist options
  • Save lychees/eb8d373e67046c6e6d6248e1be721dd3 to your computer and use it in GitHub Desktop.
Save lychees/eb8d373e67046c6e6d6248e1be721dd3 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
poly_gen.cpp
#include <lastweapon/io>
#include <assert.h>
#include <stdio.h>
#include <algorithm>
#include <iostream>
#include <vector>
using std::max;
using std::min;
using std::vector;
template <int M_> struct ModInt {
static constexpr int M = M_;
int x;
constexpr ModInt() : x(0) {}
constexpr ModInt(long long x_) : x(x_ % M) { if (x < 0) x += M; }
ModInt &operator+=(const ModInt &a) { x += a.x; if (x >= M) x -= M; return *this; }
ModInt &operator-=(const ModInt &a) { x -= a.x; if (x < 0) x += M; return *this; }
ModInt &operator*=(const ModInt &a) { x = static_cast<int>((static_cast<long long>(x) * a.x) % M); return *this; }
ModInt operator+(const ModInt &a) const { return (ModInt(*this) += a); }
ModInt operator-(const ModInt &a) const { return (ModInt(*this) -= a); }
ModInt operator*(const ModInt &a) const { return (ModInt(*this) *= a); }
ModInt operator-() const { return ModInt(-x); }
ModInt pow(long long e) const {
ModInt x2 = x, xe = 1;
for (; e; e >>= 1) {
if (e & 1) xe *= x2;
x2 *= x2;
}
return xe;
}
ModInt inv() const {
int a = x, b = M, y = 1, z = 0, t;
for (; ; ) {
t = a / b; a -= t * b;
if (a == 0) {
assert(b == 1 || b == -1);
return ModInt(b * z);
}
y -= t * z;
t = b / a; b -= t * a;
if (b == 0) {
assert(a == 1 || a == -1);
return ModInt(a * y);
}
z -= t * y;
}
}
friend ModInt operator+(long long a, const ModInt &b) { return (ModInt(a) += b); }
friend ModInt operator-(long long a, const ModInt &b) { return (ModInt(a) -= b); }
friend ModInt operator*(long long a, const ModInt &b) { return (ModInt(a) *= b); }
friend std::ostream &operator<<(std::ostream &os, const ModInt &a) { return os << a.x; }
};
/*
constexpr unsigned MO = 998244353U;
constexpr unsigned MO2 = 2U * MO;
constexpr int FFT_MAX = 23;
using Mint = ModInt<MO>;
constexpr Mint FFT_ROOTS[FFT_MAX + 1] = {1U, 998244352U, 911660635U, 372528824U, 929031873U, 452798380U, 922799308U, 781712469U, 476477967U, 166035806U, 258648936U, 584193783U, 63912897U, 350007156U, 666702199U, 968855178U, 629671588U, 24514907U, 996173970U, 363395222U, 565042129U, 733596141U, 267099868U, 15311432U};
constexpr Mint INV_FFT_ROOTS[FFT_MAX + 1] = {1U, 998244352U, 86583718U, 509520358U, 337190230U, 87557064U, 609441965U, 135236158U, 304459705U, 685443576U, 381598368U, 335559352U, 129292727U, 358024708U, 814576206U, 708402881U, 283043518U, 3707709U, 121392023U, 704923114U, 950391366U, 428961804U, 382752275U, 469870224U};
constexpr Mint FFT_RATIOS[FFT_MAX] = {911660635U, 509520358U, 369330050U, 332049552U, 983190778U, 123842337U, 238493703U, 975955924U, 603855026U, 856644456U, 131300601U, 842657263U, 730768835U, 942482514U, 806263778U, 151565301U, 510815449U, 503497456U, 743006876U, 741047443U, 56250497U, 867605899U};
constexpr Mint INV_FFT_RATIOS[FFT_MAX] = {86583718U, 372528824U, 373294451U, 645684063U, 112220581U, 692852209U, 155456985U, 797128860U, 90816748U, 860285882U, 927414960U, 354738543U, 109331171U, 293255632U, 535113200U, 308540755U, 121186627U, 608385704U, 438932459U, 359477183U, 824071951U, 103369235U};
*/
// M: prime, G: primitive root
template <int M, int G, int K> struct Fft {
using Mint = ModInt<M>;
// 1, 1/4, 1/8, 3/8, 1/16, 5/16, 3/16, 7/16, ...
Mint g[1 << (K - 1)];
Mint d[K];
constexpr Fft() : g() {
static_assert(K >= 2, "Fft: K >= 2 must hold");
static_assert(!((M - 1) & ((1 << K) - 1)), "Fft: 2^K | M - 1 must hold");
g[0] = 1;
g[1 << (K - 2)] = Mint(G).pow((M - 1) >> K);
for (int l = 1 << (K - 2); l >= 2; l >>= 1) {
g[l >> 1] = g[l] * g[l];
}
assert((g[1] * g[1]).x == M - 1);
for (int l = 2; l <= 1 << (K - 2); l <<= 1) {
for (int i = 1; i < l; ++i) {
g[l + i] = g[l] * g[i];
//g[l] * g[i+1]
}
}
d[0] = 1;
FOR(i, 1, K-1) d[i] = d[i-1] * g[1<<(i-1)].inv();
//cout << d[0] << endl;
REP(i, K-1) d[i] *= g[1<<i];
/*REP(i, K) {
cout << g[1<<i] << " ";
}
cout << endl;
cout << endl;
REP(i, K) {
cout << d[i] << " ";
}
cout << endl;
cout << M <<" " << K << endl;*/
printf("constexpr unsigned MO = %dU;\n", M);
printf("constexpr unsigned MO2 = 2U * MO;\n");
printf("constexpr int FFT_MAX = %d;\n", K);
puts("using Mint = ModInt<MO>;");
printf("constexpr Mint FFT_ROOTS[FFT_MAX + 1] = {1U");
printf(", %dU", M-1);
REP(i, K-1) {
printf(", %dU", g[1<<i]);
}
puts("};");
printf("constexpr Mint INV_FFT_ROOTS[FFT_MAX + 1] = {1U");
printf(", %dU", M-1);
REP(i, K-1) {
printf(", %dU", g[1<<i].inv());
}
puts("};");
printf("constexpr Mint FFT_RATIOS[FFT_MAX] = {%dU", d[0]);
FOR(i, 1, K-1) {
printf(", %dU", d[i]);
}
puts("};");
printf("constexpr Mint INV_FFT_RATIOS[FFT_MAX] = {%dU", d[0].inv());
FOR(i, 1, K-1) {
printf(", %dU", d[i].inv());
}
puts("};");
}
void fft(vector<Mint> &x) const {
const int n = x.size();
assert(!(n & (n - 1)) && n <= 1 << K);
for (int l = n; l >>= 1; ) {
for (int i = 0; i < (n >> 1) / l; ++i) {
for (int j = (i << 1) * l; j < (i << 1 | 1) * l; ++j) {
const Mint t = g[i] * x[j + l];
x[j + l] = x[j] - t;
x[j] += t;
}
}
}
for (int i = 0, j = 0; i < n; ++i) {
if (i < j) std::swap(x[i], x[j]);
for (int l = n; (l >>= 1) && !((j ^= l) & l); ) {}
}
}
// g[i] / g[i-1] = ratio[ctz(i)]
//
};
// 1000
// 0111
constexpr int MO = 1004535809;
constexpr int LIM = (1 << 20) + 1;
const Fft<MO, 3, 21> FFT;
using Mint = ModInt<MO>;
Mint inv[LIM];
struct ModIntPreparator {
ModIntPreparator() {
inv[1] = 1;
for (int i = 2; i < LIM; ++i) inv[i] = -(MO / i) * inv[MO % i];
}
} preparator;
struct Poly : public vector<Mint> {
Poly() {}
explicit Poly(int n) : vector<Mint>(n) {}
Poly(const vector<Mint> &vec) : vector<Mint>(vec) {}
Poly(std::initializer_list<Mint> il) : vector<Mint>(il) {}
int size() const { return vector<Mint>::size(); }
Poly take(int n) const { return Poly(vector<Mint>(this->begin(), this->begin() + min(max(n, 1), size()))); }
friend std::ostream &operator<<(std::ostream &os, const Poly &f) {
os << "[";
for (int i = 0; i < f.size(); ++i) {
if (i > 0) os << ", ";
os << f[i];
}
return os << "]";
}
Poly &operator+=(const Poly &f) {
if (size() < f.size()) resize(f.size());
for (int i = 0; i < f.size(); ++i) (*this)[i] += f[i];
return *this;
}
Poly &operator-=(const Poly &f) {
if (size() < f.size()) resize(f.size());
for (int i = 0; i < f.size(); ++i) (*this)[i] -= f[i];
return *this;
}
Poly &operator*=(const Poly &f) {
const int nt = size(), nf = f.size();
int nn;
for (nn = 1; nn < nt + nf - 1; nn <<= 1) {}
Poly ff = f;
resize(nn);
ff.resize(nn);
FFT.fft(*this);
FFT.fft(ff);
for (int i = 0; i < nn; ++i) (*this)[i] *= ff[i] * inv[nn];
std::reverse(begin() + 1, end());
FFT.fft(*this);
resize(nt + nf - 1);
return *this;
}
Poly &operator*=(const Mint &a) {
for (int i = 0; i < size(); ++i) (*this)[i] *= a;
return *this;
}
Poly operator+(const Poly &f) const { return (Poly(*this) += f); }
Poly operator-(const Poly &f) const { return (Poly(*this) -= f); }
Poly operator*(const Poly &f) const { return (Poly(*this) *= f); }
Poly operator*(const Mint &a) const { return (Poly(*this) *= a); }
friend Poly operator*(const Mint &a, const Poly &f) { return f * a; }
Poly square(int n) const {
const int nt = min(n, size());
int nn;
for (nn = 1; nn < nt + nt - 1; nn <<= 1) {}
Poly f = *this;
f.resize(nn);
FFT.fft(f);
for (int i = 0; i < nn; ++i) f[i] *= f[i] * inv[nn];
std::reverse(f.begin() + 1, f.end());
FFT.fft(f);
f.resize(n);
return f;
}
Poly differential() const {
Poly f(max(size() - 1, 1));
for (int i = 1; i < size(); ++i) f[i - 1] = i * (*this)[i];
return f;
}
Poly integral() const {
Poly f(size() + 1);
for (int i = 0; i < size(); ++i) f[i + 1] = inv[i + 1] * (*this)[i];
return f;
}
Poly exp(int n) const {
assert((*this)[0].x == 0);
const Poly d = differential();
Poly f = {1}, g = {1};
for (int m = 1; m < n; m <<= 1) {
g = g + g - (f * g.square(m)).take(m);
Poly h = d.take(m - 1);
h += (g * (f.differential() - f * h)).take(2 * m - 1);
f += (f * (take(2 * m) - h.integral())).take(2 * m);
}
return f.take(n);
}
};
// !!!Modify LIM and FFT!!!
// https://judge.yosupo.jp/problem/convolution_mod
void yosupo_convolution_mod() {
int L, M;
for (; ~scanf("%d%d", &L, &M); ) {
Poly A(L), B(M);
for (int i = 0; i < L; ++i) {
scanf("%d", &A[i].x);
}
for (int i = 0; i < M; ++i) {
scanf("%d", &B[i].x);
}
const Poly ans = A * B;
for (int i = 0; i < ans.size(); ++i) {
if (i > 0) printf(" ");
printf("%d", ans[i].x);
}
puts("");
}
}
// https://judge.yosupo.jp/problem/exp_of_formal_power_series
void yosupo_exp_of_formal_power_series() {
int N;
for (; ~scanf("%d", &N); ) {
Poly A(N);
for (int i = 0; i < N; ++i) {
scanf("%d", &A[i].x);
}
const Poly ans = A.exp(N);
for (int i = 0; i < ans.size(); ++i) {
if (i > 0) printf(" ");
printf("%d", ans[i].x);
}
puts("");
}
}
int main() {
// yosupo_convolution_mod();
yosupo_exp_of_formal_power_series();
return 0;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment