cp-library
競技プログラミング用 C++ ライブラリ
test/ntt.test.cpp
test/ntt.test.cpp · GitHub で見る
依存ライブラリ
コード
#define PROBLEM "https://judge.yosupo.jp/problem/convolution_mod"
#include "../components/ntt.hpp"
#include <bit>
#include <iostream>
#include <vector>
int main() {
std::cin.tie(nullptr)->sync_with_stdio(false);
int n, m;
std::cin >> n >> m;
int length = n + m - 1, size = std::bit_ceil(unsigned(length));
std::vector<fps_mint> a(size), b(size);
for (int i = 0; i < n; ++i) std::cin >> a[i];
for (int i = 0; i < m; ++i) std::cin >> b[i];
ntt(a);
ntt(b);
for (int i = 0; i < size; ++i) a[i] *= b[i];
ntt(a, true);
for (int i = 0; i < length; ++i) std::cout << a[i] << ' ';
std::cout << '\n';
}
#line 1 "test/ntt.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/convolution_mod"
#line 1 "components/ntt.hpp"
#include <algorithm>
#include <cassert>
#include <cstddef>
#include <initializer_list>
#include <memory>
#include <new>
#include <type_traits>
#include <utility>
#include <vector>
#line 1 "components/mint_arithmetic.hpp"
#line 1 "math/mint.hpp"
#line 4 "math/mint.hpp"
#include <istream>
#include <ostream>
#line 7 "math/mint.hpp"
#if defined(__GNUC__) && !defined(__clang__) && (defined(__x86_64__) || defined(__i386__))
#pragma GCC push_options
#pragma GCC optimize("O3")
#endif
namespace mint_detail {
template <class I>
constexpr bool mint_integral = std::is_integral_v<I>
#if defined(__SIZEOF_INT128__)
|| std::is_same_v<I, __int128> || std::is_same_v<I, unsigned __int128>
#endif
;
struct [[gnu::warn_unused]] Mint {
using mint = Mint;
unsigned v = 0;
[[nodiscard]] static constexpr unsigned mod() { return 998244353; }
[[nodiscard, gnu::always_inline]] static unsigned reduce(unsigned long long x) {
return (x + (unsigned long long)(unsigned(x) * 998244351u) * mod()) >> 32;
}
[[gnu::always_inline]] Mint() = default;
template <class I, std::enable_if_t<mint_integral<I>, int> = 0>
Mint(I x) {
long long r = x % (long long)mod();
if (r < 0) r += mod();
v = reduce((unsigned long long)r * 932051910);
}
[[nodiscard, gnu::always_inline]] static Mint raw(unsigned x) {
Mint a;
a.v = reduce((unsigned long long)x * 932051910);
return a;
}
[[nodiscard, gnu::always_inline]] static Mint word(unsigned x) {
Mint a;
a.v = x;
return a;
}
[[nodiscard, gnu::always_inline]] unsigned val() const {
unsigned r = reduce(v);
return r < mod() ? r : r - mod();
}
friend std::istream &operator>>(std::istream &is, Mint &x) {
long long value;
if (is >> value) x = Mint(value);
return is;
}
friend std::ostream &operator<<(std::ostream &os, const Mint &x) {
return os << x.val();
}
[[gnu::always_inline]] Mint &operator+=(Mint b) {
v += b.v;
if (v >= 2 * mod()) v -= 2 * mod();
return *this;
}
[[gnu::always_inline]] Mint &operator-=(Mint b) {
v -= b.v;
if (v >> 31) v += 2 * mod();
return *this;
}
[[gnu::always_inline]] Mint &operator*=(Mint b) {
v = reduce((unsigned long long)v * b.v);
return *this;
}
Mint &operator/=(Mint b) { return *this *= b.inv(); }
Mint &operator++() { return *this += Mint(1); }
Mint &operator--() { return *this -= Mint(1); }
Mint operator++(int) {
Mint a = *this;
++*this;
return a;
}
Mint operator--(int) {
Mint a = *this;
--*this;
return a;
}
[[nodiscard, gnu::always_inline]] Mint operator+() const { return *this; }
[[nodiscard, gnu::always_inline]] Mint operator-() const { return word(v ? 2 * mod() - v : 0); }
[[nodiscard, gnu::always_inline]] Mint operator+(Mint b) const { return Mint(*this) += b; }
[[nodiscard, gnu::always_inline]] Mint operator-(Mint b) const { return Mint(*this) -= b; }
[[nodiscard, gnu::always_inline]] Mint operator*(Mint b) const { return Mint(*this) *= b; }
[[nodiscard]] Mint operator/(Mint b) const { return Mint(*this) /= b; }
[[nodiscard, gnu::always_inline]] bool operator==(Mint b) const { return (v < mod() ? v : v - mod()) == (b.v < mod() ? b.v : b.v - mod()); }
[[nodiscard, gnu::always_inline]] bool operator!=(Mint b) const { return !(*this == b); }
template <class I, std::enable_if_t<mint_integral<I>, int> = 0>
[[nodiscard]] friend Mint operator+(I a, Mint b) { return Mint(a) + b; }
template <class I, std::enable_if_t<mint_integral<I>, int> = 0>
[[nodiscard]] friend Mint operator-(I a, Mint b) { return Mint(a) - b; }
template <class I, std::enable_if_t<mint_integral<I>, int> = 0>
[[nodiscard]] friend Mint operator*(I a, Mint b) { return Mint(a) * b; }
template <class I, std::enable_if_t<mint_integral<I>, int> = 0>
[[nodiscard]] friend Mint operator/(I a, Mint b) { return Mint(a) / b; }
template <class I, std::enable_if_t<mint_integral<I>, int> = 0>
[[nodiscard]] friend bool operator==(I a, Mint b) { return Mint(a) == b; }
template <class I, std::enable_if_t<mint_integral<I>, int> = 0>
[[nodiscard]] friend bool operator!=(I a, Mint b) { return Mint(a) != b; }
[[nodiscard]] Mint pow(long long n) const {
assert(n >= 0);
if (!n) return Mint(1);
Mint r(1), a = *this;
for (; n > 1; n >>= 1, a *= a)
if (n & 1) r *= a;
return r * a;
}
[[nodiscard, gnu::noinline]] Mint inv() const {
assert(*this != Mint(0));
Mint a2 = *this * *this, a3 = a2 * *this;
Mint a7 = a3 * a3 * *this;
Mint a31 = a7;
for (int i = 0; i < 2; ++i) a31 *= a31;
a31 *= a3;
Mint r = a7;
for (int i = 0; i < 3; ++i) r *= r;
r *= a3;
r *= r;
for (int j = 0; j < 4; ++j) {
for (int i = 0; i < 5; ++i) r *= r;
r *= a31;
}
for (int i = 0; i < 3; ++i) r *= r;
return r * a7;
}
};
} // namespace mint_detail
using mint = mint_detail::Mint;
using fps_mint = mint;
#if defined(__GNUC__) && !defined(__clang__) && (defined(__x86_64__) || defined(__i386__))
#pragma GCC pop_options
#endif
#line 5 "components/mint_arithmetic.hpp"
#if defined(__x86_64__) || defined(__i386__)
#include <immintrin.h>
#endif
#if defined(__GNUC__) && !defined(__clang__) && (defined(__x86_64__) || defined(__i386__))
#pragma GCC push_options
#pragma GCC optimize("O3")
#endif
namespace mint_detail {
struct Ops {
using T = mint;
typedef unsigned int U __attribute__((may_alias));
static constexpr U mod = 998244353, twice = 2 * mod, one = 301989884;
static U reduce(unsigned long long x) { return T::reduce(x); }
static U mul(U a, U b) { return reduce((unsigned long long)a * b); }
static U add(U a, U b) {
U c = a + b;
return c < twice ? c : c - twice;
}
static U sub(U a, U b) {
U c = a - b;
return c < twice ? c : c + twice;
}
static U in(U a) { return mul(a, 932051910); }
static U out(U a) {
a = reduce(a);
return a < mod ? a : a - mod;
}
static U quotient(U r) { return ((unsigned long long)r << 32) / mod; }
static U fixed_mul(U a, U r, U q) {
return a * r - U(((unsigned long long)a * q) >> 32) * mod;
}
#if defined(__x86_64__) || defined(__i386__)
static bool vectorized() {
#if defined(__AVX2__)
return true;
#else
static const bool available = __builtin_cpu_supports("avx2");
return available;
#endif
}
using I8 = __m256i;
__attribute__((target("avx2"))) static I8 load8(const U *a) {
return _mm256_loadu_si256(reinterpret_cast<const I8 *>(a));
}
__attribute__((target("avx2"))) static void store8(U *a, I8 x) {
_mm256_storeu_si256(reinterpret_cast<I8 *>(a), x);
}
__attribute__((target("avx2"), always_inline)) static inline I8 reduce8(I8 x, I8 y) {
I8 r = _mm256_set1_epi32(998244351), p = _mm256_set1_epi32(mod);
x = _mm256_add_epi64(x, _mm256_mul_epu32(_mm256_mul_epu32(x, r), p));
y = _mm256_add_epi64(y, _mm256_mul_epu32(_mm256_mul_epu32(y, r), p));
return _mm256_blend_epi32(_mm256_srli_epi64(x, 32), y, 0xaa);
}
__attribute__((target("avx2"))) static I8 mul8(I8 a, I8 b) {
return reduce8(_mm256_mul_epu32(a, b), _mm256_mul_epu32(_mm256_srli_epi64(a, 32), _mm256_srli_epi64(b, 32)));
}
__attribute__((target("avx2"))) static I8 fixed_mul8(I8 a, I8 r, I8 q) {
I8 x = _mm256_srli_epi64(_mm256_mul_epu32(a, q), 32);
I8 y = _mm256_mul_epu32(_mm256_srli_epi64(a, 32), q);
I8 t = _mm256_blend_epi32(x, y, 0xaa);
return _mm256_sub_epi32(_mm256_mullo_epi32(a, r),
_mm256_mullo_epi32(t, _mm256_set1_epi32(mod)));
}
__attribute__((target("avx2"))) static I8 fixed_even8(I8 a, I8 r, I8 q) {
I8 t = _mm256_srli_epi64(_mm256_mul_epu32(a, q), 32);
return _mm256_sub_epi64(_mm256_mul_epu32(a, r),
_mm256_mul_epu32(t, _mm256_set1_epi32(mod)));
}
__attribute__((target("avx2"))) static I8 add8(I8 a, I8 b) {
I8 x = _mm256_add_epi32(a, b);
return _mm256_min_epu32(x, _mm256_sub_epi32(x, _mm256_set1_epi32(twice)));
}
__attribute__((target("avx2"))) static I8 sub8(I8 a, I8 b) {
I8 x = _mm256_sub_epi32(a, b);
return _mm256_min_epu32(x, _mm256_add_epi32(x, _mm256_set1_epi32(twice)));
}
#endif
};
}
#if defined(__GNUC__) && !defined(__clang__) && (defined(__x86_64__) || defined(__i386__))
#pragma GCC pop_options
#endif
#line 15 "components/ntt.hpp"
#if defined(__GNUC__) && !defined(__clang__) && (defined(__x86_64__) || defined(__i386__))
#pragma GCC push_options
#pragma GCC optimize("O3")
#endif
namespace fps_detail {
struct NTT : mint_detail::Ops {
using V = std::vector<unsigned int>;
V roots, inverse_roots;
V root_quotients, inverse_root_quotients;
mutable V cube_roots, inverse_cube_roots, cube_quotients, inverse_cube_quotients;
void prepare_cubes(int n) const {
int begin = cube_roots.size();
if (begin >= n) return;
cube_roots.resize(n);
inverse_cube_roots.resize(n);
cube_quotients.resize(n);
inverse_cube_quotients.resize(n);
for (int k = begin; k < n; ++k) {
U x = (unsigned long long)roots[k] * roots[2 * k] % mod;
U y = (unsigned long long)inverse_roots[k] * inverse_roots[2 * k] % mod;
cube_roots[k] = x;
inverse_cube_roots[k] = y;
cube_quotients[k] = quotient(x);
inverse_cube_quotients[k] = quotient(y);
}
}
NTT(int size, [[maybe_unused]] bool composition = true)
: roots(size, one),
inverse_roots(size, one),
root_quotients(size),
inverse_root_quotients(size) {
for (int h = 1; h < size; h <<= 1) {
U r = in(T(3).pow((mod - 1) / (4 * h)).val());
U ir = in(T(3).pow(mod - 1 - (mod - 1) / (4 * h)).val());
for (int i = 0; i < h; ++i) {
roots[h + i] = mul(roots[i], r);
inverse_roots[h + i] = mul(inverse_roots[i], ir);
}
}
for (int i = 0; i < size; ++i) {
root_quotients[i] = (roots[i] < mod ? roots[i] : roots[i] - mod) * 998244351u;
inverse_root_quotients[i] = (inverse_roots[i] < mod ? inverse_roots[i] : inverse_roots[i] - mod) * 998244351u;
roots[i] = out(roots[i]);
inverse_roots[i] = out(inverse_roots[i]);
}
}
#if defined(__x86_64__) || defined(__i386__)
template <int stride>
__attribute__((target("avx2"))) static I8 rates8(const V &values, int k) {
if constexpr (stride == 1)
return _mm256_cvtepu32_epi64(
_mm_loadu_si128(reinterpret_cast<const __m128i *>(values.data() + 4 * k)));
if constexpr (stride == 2)
return _mm256_permute4x64_epi64(
_mm256_cvtepu32_epi64(
_mm_loadl_epi64(reinterpret_cast<const __m128i *>(values.data() + 2 * k))),
0x50);
if constexpr (stride == 4) return _mm256_set1_epi32(values[k]);
}
template <int stride, bool inverse>
__attribute__((target("avx2"))) I8 small8(I8 z, int k) const {
if constexpr (stride == 4) {
if (k == 0) {
I8 x = _mm256_permute2x128_si256(z, z, 0x00),
y = _mm256_permute2x128_si256(z, z, 0x11);
return _mm256_blend_epi32(add8(x, y), sub8(x, y), 0xf0);
}
}
const V &w = inverse ? inverse_roots : roots;
const V " = inverse ? inverse_root_quotients : root_quotients;
I8 r = rates8<stride>(w, k), q = rates8<stride>(quot, k);
if constexpr (inverse) {
if constexpr (stride == 2) z = _mm256_shuffle_epi32(z, 0xd8);
if constexpr (stride == 4)
z = _mm256_permutevar8x32_epi32(z, _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7));
I8 y = _mm256_srli_epi64(z, 32), x = add8(z, y);
I8 d = fixed_even8(sub8(z, y), r, q);
z = _mm256_blend_epi32(x, _mm256_slli_epi64(d, 32), 0xaa);
} else {
if constexpr (stride == 1) z = _mm256_shuffle_epi32(z, 0xb1);
if constexpr (stride == 2) z = _mm256_shuffle_epi32(z, 0x72);
if constexpr (stride == 4)
z = _mm256_permutevar8x32_epi32(z, _mm256_setr_epi32(4, 0, 5, 1, 6, 2, 7, 3));
I8 y = fixed_even8(z, r, q);
I8 x = _mm256_blend_epi32(_mm256_srli_epi64(z, 32), z, 0xaa);
y = _mm256_blend_epi32(
y, _mm256_sub_epi32(_mm256_set1_epi32(twice), _mm256_slli_epi64(y, 32)), 0xaa);
z = add8(x, y);
}
if constexpr (stride == 2) z = _mm256_shuffle_epi32(z, 0xd8);
if constexpr (stride == 4)
z = _mm256_permutevar8x32_epi32(z, _mm256_setr_epi32(0, 2, 4, 6, 1, 3, 5, 7));
return z;
}
template <bool inverse>
__attribute__((target("avx2"))) void fft8(U *a, int length, int limit = 0) const {
if (!limit) limit = length;
if constexpr (!inverse) {
int h = limit / 2;
for (; h >= 16; h >>= 2) radix4_8<false>(a, h / 2, length);
if (h == 8) {
if (limit >= 32)
for (int i = 0; i < length; i += 8) {
I8 x = load8(a + i), p = _mm256_set1_epi32(twice);
store8(a + i, _mm256_min_epu32(x, _mm256_sub_epi32(x, p)));
}
stage8<false>(a, h, length);
}
}
for (int j = 0; j < length; j += 8) {
I8 z = load8(a + j);
if constexpr (!inverse)
if ((limit >= 32) && (__builtin_ctz((unsigned int)limit) & 1)) {
I8 p = _mm256_set1_epi32(twice);
z = _mm256_min_epu32(z, _mm256_sub_epi32(z, p));
}
if constexpr (inverse)
z = small8<4, true>(small8<2, true>(small8<1, true>(z, j / 8), j / 8), j / 8);
else
z = small8<1, false>(small8<2, false>(small8<4, false>(z, j / 8), j / 8),
j / 8);
store8(a + j, z);
}
if constexpr (inverse) {
int h = 8;
for (; 4 * h <= limit; h <<= 2) radix4_8<true>(a, h, length);
if (h < limit) stage8<true>(a, h, length);
}
}
template <bool inverse>
__attribute__((target("avx2"))) void radix4_8(U *a, int stride, int length) const {
const V &w = inverse ? inverse_roots : roots;
const V &q = inverse ? inverse_root_quotients : root_quotients;
prepare_cubes(length / (4 * stride));
const V &w3 = inverse ? inverse_cube_roots : cube_roots;
const V &q3s = inverse ? inverse_cube_quotients : cube_quotients;
I8 p = _mm256_set1_epi32(twice);
I8 im = _mm256_set1_epi32(w[1]), iq = _mm256_set1_epi32(q[1]);
for (int j = 0, k = 0; j < length; j += 4 * stride, ++k) {
I8 r = _mm256_set1_epi32(w[k]), rq = _mm256_set1_epi32(q[k]);
I8 r0 = _mm256_set1_epi32(w[2 * k]), q0 = _mm256_set1_epi32(q[2 * k]);
I8 r3 = _mm256_set1_epi32(w3[k]), q3 = _mm256_set1_epi32(q3s[k]);
for (int i = 0; i < stride; i += 8) {
I8 x0 = load8(a + j + i), x1 = load8(a + j + stride + i);
I8 x2 = load8(a + j + 2 * stride + i), x3 = load8(a + j + 3 * stride + i);
if constexpr (inverse) {
I8 u = add8(x0, x1), v = sub8(x0, x1);
I8 s = add8(x2, x3);
I8 t = fixed_mul8(_mm256_sub_epi32(_mm256_add_epi32(x2, p), x3), im, iq);
x0 = add8(u, s);
x1 = _mm256_add_epi32(v, t);
x2 = _mm256_sub_epi32(_mm256_add_epi32(u, p), s);
x3 = _mm256_sub_epi32(_mm256_add_epi32(v, p), t);
if (k) {
x1 = fixed_mul8(x1, r0, q0);
x2 = fixed_mul8(x2, r, rq);
x3 = fixed_mul8(x3, r3, q3);
} else {
x1 = _mm256_min_epu32(x1, _mm256_sub_epi32(x1, p));
x2 = _mm256_min_epu32(x2, _mm256_sub_epi32(x2, p));
x3 = _mm256_min_epu32(x3, _mm256_sub_epi32(x3, p));
}
} else {
x0 = _mm256_min_epu32(x0, _mm256_sub_epi32(x0, p));
if (!k) {
x1 = _mm256_min_epu32(x1, _mm256_sub_epi32(x1, p));
x2 = _mm256_min_epu32(x2, _mm256_sub_epi32(x2, p));
x3 = _mm256_min_epu32(x3, _mm256_sub_epi32(x3, p));
}
if (k) {
x1 = fixed_mul8(x1, r0, q0);
x2 = fixed_mul8(x2, r, rq);
x3 = fixed_mul8(x3, r3, q3);
}
I8 u = add8(x0, x2), v = sub8(x0, x2), s = add8(x1, x3);
I8 t = fixed_mul8(_mm256_sub_epi32(_mm256_add_epi32(x1, p), x3), im, iq);
x0 = _mm256_add_epi32(u, s);
x1 = _mm256_sub_epi32(_mm256_add_epi32(u, p), s);
x2 = _mm256_add_epi32(v, t);
x3 = _mm256_sub_epi32(_mm256_add_epi32(v, p), t);
}
store8(a + j + i, x0);
store8(a + j + stride + i, x1);
store8(a + j + 2 * stride + i, x2);
store8(a + j + 3 * stride + i, x3);
}
}
}
template <bool inverse>
__attribute__((target("avx2"))) void stage8(U *a, int stride, int length, int first = 0) const {
const V &w = inverse ? inverse_roots : roots;
const V " = inverse ? inverse_root_quotients : root_quotients;
if (stride < 8) {
for (int j = 0; j < length; j += 8) {
I8 z = load8(a + j);
if (stride == 1)
z = small8<1, inverse>(z, j / 8);
else if (stride == 2)
z = small8<2, inverse>(z, j / 8);
else
z = small8<4, inverse>(z, j / 8);
store8(a + j, z);
}
} else {
for (int j = 0, k = first; j < length; j += 2 * stride, ++k) {
I8 r = _mm256_set1_epi32(w[k]), q = _mm256_set1_epi32(quot[k]);
for (int i = 0; i < stride; i += 8) {
I8 x = load8(a + j + i), y = load8(a + j + stride + i);
if constexpr (!inverse)
if (k) y = fixed_mul8(y, r, q);
I8 z = sub8(x, y);
x = add8(x, y);
if constexpr (inverse)
if (k) z = fixed_mul8(z, r, q);
store8(a + j + i, x);
store8(a + j + stride + i, z);
}
}
}
}
__attribute__((target("avx2"))) static void load_pairs(const U *a, I8 &x, I8 &y) {
I8 index = _mm256_setr_epi32(0, 2, 4, 6, 1, 3, 5, 7);
I8 u = _mm256_permutevar8x32_epi32(load8(a), index),
v = _mm256_permutevar8x32_epi32(load8(a + 8), index);
x = _mm256_permute2x128_si256(u, v, 0x20);
y = _mm256_permute2x128_si256(u, v, 0x31);
}
__attribute__((target("avx2"))) static void scale_copy8(U *dst, const U *src, int length, U r) {
U quotient_r = quotient(r);
I8 v = _mm256_set1_epi32(r), q = _mm256_set1_epi32(quotient_r);
int i = 0;
for (; i + 8 <= length; i += 8) store8(dst + i, fixed_mul8(load8(src + i), v, q));
for (; i < length; ++i) dst[i] = fixed_mul(src[i], r, quotient_r);
}
__attribute__((target("avx2"))) static void scale8(U *a, int length, U r) {
scale_copy8(a, a, length, r);
}
#endif
template <bool inverse>
void stage(U *a, int stride, int length, int first = 0) const {
const V &w = inverse ? inverse_roots : roots;
const V " = inverse ? inverse_root_quotients : root_quotients;
#if defined(__x86_64__) || defined(__i386__)
if (length >= 8 && vectorized()) {
stage8<inverse>(a, stride, length, first);
return;
}
#endif
for (int j = 0, k = first; j < length; j += 2 * stride, ++k) {
U r = w[k], q = quot[k];
int i = 0;
for (; i < stride; ++i) {
U x = a[j + i], y = a[j + stride + i];
if constexpr (!inverse)
if (k) y = fixed_mul(y, r, q);
a[j + i] = add(x, y);
a[j + stride + i] = inverse && k ? fixed_mul(sub(x, y), r, q) : sub(x, y);
}
}
}
template <bool inverse>
void fft(U *a, int length, int limit = 0) const {
if (!limit) limit = length;
#if defined(__x86_64__) || defined(__i386__)
if (limit >= 8 && vectorized()) {
fft8<inverse>(a, length, limit);
return;
}
#endif
if constexpr (inverse) {
for (int h = 1; h < limit; h <<= 1) stage<true>(a, h, length);
} else {
for (int h = limit / 2; h; h >>= 1) stage<false>(a, h, length);
}
}
#if defined(__x86_64__) || defined(__i386__)
__attribute__((target("avx2"))) static void product8(U *dst, const U *a, const U *b, int n) {
int i = 0;
for (; i + 8 <= n; i += 8) {
I8 x = mul8(load8(a + i), load8(b + i));
store8(dst + i, x);
}
for (; i < n; ++i) dst[i] = mul(a[i], b[i]);
}
__attribute__((target("avx2"))) static void weighted8(U *dst, const U *src, int n, int first) {
alignas(32) U factors[8];
for (int j = 0; j < 8; ++j) factors[j] = in(first + j);
I8 r = load8(factors), step = _mm256_set1_epi32(in(8));
int i = 0;
for (; i + 8 <= n; i += 8) {
store8(dst + i, mul8(load8(src + i), r));
r = add8(r, step);
}
for (; i < n; ++i) dst[i] = mul(src[i], in(first + i));
}
__attribute__((target("avx2"), always_inline)) static I8 inverse8(I8 a) {
I8 a2 = mul8(a, a), a3 = mul8(a2, a);
I8 a7 = mul8(mul8(a3, a3), a), a31 = a7;
for (int repeat_index = 0; repeat_index < 2; ++repeat_index) a31 = mul8(a31, a31);
a31 = mul8(a31, a3);
I8 r = a7;
for (int repeat_index = 0; repeat_index < 3; ++repeat_index) r = mul8(r, r);
r = mul8(r, a3);
r = mul8(r, r);
for (int j = 0; j < 4; ++j) {
for (int repeat_index = 0; repeat_index < 5; ++repeat_index) r = mul8(r, r);
r = mul8(r, a31);
}
for (int repeat_index = 0; repeat_index < 3; ++repeat_index) r = mul8(r, r);
return mul8(r, a7);
}
__attribute__((target("avx2"))) static void batch_inverse8(U *dst, const U *src, int n) {
if (n < 16) {
for (int i = 0; i < n; ++i) dst[i] = T::word(src[i]).inv().v;
return;
}
int count = n / 16 * 16;
I8 p = load8(src), q = load8(src + 8);
store8(dst, p);
store8(dst + 8, q);
for (int i = 16; i < count; i += 16) {
p = mul8(p, load8(src + i));
q = mul8(q, load8(src + i + 8));
store8(dst + i, p);
store8(dst + i + 8, q);
}
alignas(32) U lanes[16];
store8(lanes, p);
store8(lanes + 8, q);
for ([[maybe_unused]] auto x : lanes) assert(out(x) != 0);
I8 inverse = inverse8(p), other = inverse8(q);
for (int i = count - 16; i >= 16; i -= 16) {
store8(dst + i, mul8(inverse, load8(dst + i - 16)));
store8(dst + i + 8, mul8(other, load8(dst + i - 8)));
inverse = mul8(inverse, load8(src + i));
other = mul8(other, load8(src + i + 8));
}
store8(dst, inverse);
store8(dst + 8, other);
for (int i = count; i < n; ++i) dst[i] = T::word(src[i]).inv().v;
}
__attribute__((target("avx2"))) static void inverse_sequence8(U *dst, int first, int n) {
alignas(32) U initial[8];
for (int j = 0; j < 8; ++j) initial[j] = in(first + j);
I8 factor = load8(initial), step = _mm256_set1_epi32(in(8));
I8 p = factor;
store8(dst, p);
for (int i = 8; i < n; i += 8) {
factor = add8(factor, step);
p = mul8(p, factor);
store8(dst + i, p);
}
I8 inverse = inverse8(p);
for (int i = n - 8; i >= 8; i -= 8) {
store8(dst + i, mul8(inverse, load8(dst + i - 8)));
inverse = mul8(inverse, factor);
factor = sub8(factor, step);
}
store8(dst, inverse);
}
__attribute__((target("avx2"))) static void normalize8(U *a, int n) {
I8 p = _mm256_set1_epi32(mod);
int i = 0;
for (; i + 8 <= n; i += 8) {
I8 x = load8(a + i);
store8(a + i, _mm256_min_epu32(x, _mm256_sub_epi32(x, p)));
}
for (; i < n; ++i)
if (a[i] >= mod) a[i] -= mod;
}
#endif
static void weighted(U *dst, const U *src, int n, int first) {
#if defined(__x86_64__) || defined(__i386__)
if (n >= 8 && vectorized()) {
weighted8(dst, src, n, first);
return;
}
#endif
for (int i = 0; i < n; ++i) dst[i] = mul(src[i], in(first + i));
}
static void normalize(U *a, int n) {
#if defined(__x86_64__) || defined(__i386__)
if (n >= 8 && vectorized()) {
normalize8(a, n);
return;
}
#endif
int i = 0;
for (; i < n; ++i)
if (a[i] >= mod) a[i] -= mod;
}
static void product(U *dst, const U *a, const U *b, int n) {
#if defined(__x86_64__) || defined(__i386__)
if (n >= 8 && vectorized()) {
product8(dst, a, b, n);
return;
}
#endif
int i = 0;
for (; i < n; ++i) dst[i] = mul(a[i], b[i]);
}
#if defined(__x86_64__) || defined(__i386__)
__attribute__((target("avx2"), always_inline)) void partial_stage8(U *a, int stride, int k) const {
const I8 p = _mm256_set1_epi32(twice);
const I8 r = _mm256_set1_epi32(roots[k]), q = _mm256_set1_epi32(root_quotients[k]);
for (int i = 0; i < stride; i += 8) {
I8 x = load8(a + i), y = load8(a + stride + i);
x = _mm256_min_epu32(x, _mm256_sub_epi32(x, p));
y = k ? fixed_mul8(y, r, q) : _mm256_min_epu32(y, _mm256_sub_epi32(y, p));
store8(a + i, _mm256_add_epi32(x, y));
store8(a + stride + i, _mm256_sub_epi32(_mm256_add_epi32(x, p), y));
}
}
template <int count, bool second = false>
__attribute__((target("avx2"), always_inline)) void inverse_leaf_products8(U *a, U *b, const U *ws, const U *qs) const {
I8 p = _mm256_set1_epi32(mod), pp = _mm256_set1_epi32(twice);
alignas(32) U window[count][16], coefficients[count][8];
I8 e[count], o[count];
for (int k = 0; k < count; ++k) {
I8 x = load8(a + 8 * k), y = load8(b + 8 * k);
x = _mm256_min_epu32(x, _mm256_sub_epi32(x, pp));
if constexpr (!second) y = _mm256_min_epu32(y, _mm256_sub_epi32(y, pp));
x = _mm256_min_epu32(x, _mm256_sub_epi32(x, p));
if constexpr (!second) y = _mm256_min_epu32(y, _mm256_sub_epi32(y, p));
I8 low = fixed_mul8(x, _mm256_set1_epi32(ws[k]), _mm256_set1_epi32(qs[k]));
store8(window[k], _mm256_min_epu32(low, _mm256_sub_epi32(low, p)));
store8(window[k] + 8, x);
store8(coefficients[k], y);
if constexpr (!second) store8(b + 8 * k, y);
e[k] = o[k] = _mm256_setzero_si256();
}
for (int i = 0; i < 8; ++i)
for (int k = 0; k < count; ++k) {
I8 z = load8(window[k] + 8 - i), r = _mm256_set1_epi32(coefficients[k][i]);
e[k] = _mm256_add_epi64(e[k], _mm256_mul_epu32(z, r));
o[k] = _mm256_add_epi64(o[k], _mm256_mul_epu32(_mm256_srli_epi64(z, 32), r));
}
for (int k = 0; k < count; ++k) {
I8 x = reduce8(e[k], o[k]);
store8(a + 8 * k, _mm256_min_epu32(x, _mm256_sub_epi32(x, pp)));
}
}
__attribute__((target("avx2"))) static U reverse_dot8(const U *a, const U *b, int n) {
I8 p = _mm256_set1_epi32(mod), pp = _mm256_set1_epi32(twice);
I8 reversed = _mm256_setr_epi32(7, 6, 5, 4, 3, 2, 1, 0), total = _mm256_setzero_si256();
int i = 0;
for (; i + 8 <= n;) {
I8 even = _mm256_setzero_si256(), odd = even;
int end = std::min(n / 8 * 8, i + 64);
for (; i < end; i += 8) {
I8 x = load8(a + i), y = _mm256_permutevar8x32_epi32(load8(b - i - 7), reversed);
x = _mm256_min_epu32(x, _mm256_sub_epi32(x, pp));
y = _mm256_min_epu32(y, _mm256_sub_epi32(y, pp));
x = _mm256_min_epu32(x, _mm256_sub_epi32(x, p));
y = _mm256_min_epu32(y, _mm256_sub_epi32(y, p));
even = _mm256_add_epi64(even, _mm256_mul_epu32(x, y));
odd = _mm256_add_epi64(odd, _mm256_mul_epu32(_mm256_srli_epi64(x, 32), _mm256_srli_epi64(y, 32)));
}
I8 value = reduce8(even, odd);
value = _mm256_min_epu32(value, _mm256_sub_epi32(value, pp));
total = add8(total, value);
}
total = add8(total, _mm256_permute2x128_si256(total, total, 1));
total = add8(total, _mm256_shuffle_epi32(total, 0x4e));
total = add8(total, _mm256_shuffle_epi32(total, 0xb1));
U result = _mm_cvtsi128_si32(_mm256_castsi256_si128(total));
for (; i < n; ++i) {
U x = a[i], y = b[-i];
if (x >= twice) x -= twice;
if (y >= twice) y -= twice;
if (x >= mod) x -= mod;
if (y >= mod) y -= mod;
result = add(result, mul(x, y));
}
return result;
}
template <bool inverse>
__attribute__((target("avx2"))) void inverse_partial8(U *a, int n, int limit = 0) const {
if (!limit) limit = n;
if constexpr (!inverse) {
int stride = limit / 2;
for (; stride >= 16; stride >>= 2) radix4_8<false>(a, stride / 2, n);
if (stride == 8)
for (int i = 0; i < n; i += 16) partial_stage8(a + i, 8, i / 16);
} else {
int stride = 8;
for (; 4 * stride <= n; stride <<= 2) radix4_8<true>(a, stride, n);
if (stride < n) stage8<true>(a, stride, n);
}
}
__attribute__((target("avx2"))) void inverse_newton8(U *a, U *b, int n) const {
assert(n >= 64 && !(n & (n - 1)) && n <= (1 << 23) && static_cast<int>(roots.size()) >= n / 16);
std::copy_n(b, n / 2, b + n / 2);
inverse_partial8<false>(b, n, n / 2);
inverse_partial8<false>(a, n);
for (int i = 0; i < n; i += 64) {
U ws[8], qs[8];
for (int j = 0; j < 8; ++j) {
int k = i / 8 + j;
ws[j] = (k & 1) ? mod - roots[k / 2] : roots[k / 2];
qs[j] = (k & 1) ? ~root_quotients[k / 2] : root_quotients[k / 2];
}
inverse_leaf_products8<8>(a + i, b + i, ws, qs);
}
inverse_partial8<true>(a, n);
std::fill_n(a, n / 2, U(0));
inverse_partial8<false>(a, n);
for (int i = 0; i < n; i += 64) {
U ws[8], qs[8];
for (int j = 0; j < 8; ++j) {
int k = i / 8 + j;
ws[j] = (k & 1) ? mod - roots[k / 2] : roots[k / 2];
qs[j] = (k & 1) ? ~root_quotients[k / 2] : root_quotients[k / 2];
}
inverse_leaf_products8<8, true>(a + i, b + i, ws, qs);
}
inverse_partial8<true>(a, n);
U norm = mod - (mod - 1) / (n / 8);
scale8(a + n / 2, n / 2, mod - (unsigned long long)norm * norm % mod);
}
#endif
#if defined(__x86_64__) || defined(__i386__)
template <bool odd>
__attribute__((target("avx2"))) static void reduce_pairs8(U *a, U *b, const U *twist, int n) {
for (int i = 0; i < n; i += 8) {
I8 x, y, u, v;
load_pairs(a + 2 * i, x, y);
load_pairs(b + 2 * i, u, v);
x = mul8(x, v);
y = mul8(y, u);
if constexpr (odd) x = mul8(sub8(x, y), load8(twist + i));
else x = add8(x, y);
store8(a + i, x);
store8(b + i, mul8(u, v));
}
}
__attribute__((target("avx2"))) static void combine8(U *a, const U *b, const U *x, const U *y, int n) {
for (int i = 0; i < n; i += 8) store8(a + i, add8(mul8(load8(a + i), load8(y + i)), mul8(load8(b + i), load8(x + i))));
}
#endif
template <bool odd>
static void reduce_pairs(U *a, U *b, const U *twist, int n) {
#if defined(__x86_64__) || defined(__i386__)
if (n >= 8 && vectorized()) {
reduce_pairs8<odd>(a, b, twist, n);
return;
}
#endif
int i = 0;
for (; i < n; ++i) {
U u = mul(a[2 * i], b[2 * i + 1]), v = mul(a[2 * i + 1], b[2 * i]);
a[i] = odd ? mul(sub(u, v), twist[i]) : add(u, v);
b[i] = mul(b[2 * i], b[2 * i + 1]);
}
}
static void combine(U *a, const U *b, const U *x, const U *y, int n) {
#if defined(__x86_64__) || defined(__i386__)
if (n >= 8 && vectorized()) {
combine8(a, b, x, y, n);
return;
}
#endif
int i = 0;
for (; i < n; ++i) a[i] = add(mul(a[i], y[i]), mul(b[i], x[i]));
}
#if defined(__x86_64__) || defined(__i386__)
__attribute__((target("avx2"))) static void geometric8(U *a, int n, U first, U step) {
alignas(32) U factors[8];
U power = one;
for (int i = 0; i < 8; ++i) {
factors[i] = mul(first, power);
power = mul(power, step);
}
I8 r = load8(factors), ratio = _mm256_set1_epi32(power);
int i = 0;
for (; i + 8 <= n; i += 8) {
store8(a + i, mul8(load8(a + i), r));
r = mul8(r, ratio);
}
store8(factors, r);
for (int j = 0; i < n; ++i, ++j) a[i] = mul(a[i], factors[j]);
}
#endif
static void geometric(U *a, int n, U first, U step) {
#if defined(__x86_64__) || defined(__i386__)
if (n >= 8 && vectorized()) {
geometric8(a, n, first, step);
return;
}
#endif
int i = 0;
for (; i < n; ++i) {
a[i] = mul(a[i], first);
first = mul(first, step);
}
}
#if defined(__x86_64__) || defined(__i386__)
__attribute__((target("avx2"))) static void add_scaled8(U *a, const U *b, int n, U r) {
U q = quotient(r);
I8 v = _mm256_set1_epi32(r), vq = _mm256_set1_epi32(q);
int i = 0;
for (; i + 8 <= n; i += 8) store8(a + i, add8(load8(a + i), fixed_mul8(load8(b + i), v, vq)));
for (; i < n; ++i) a[i] = add(a[i], fixed_mul(b[i], r, q));
}
#endif
static void add_scaled(U *a, const U *b, int n, U r) {
#if defined(__x86_64__) || defined(__i386__)
if (n >= 8 && vectorized()) {
add_scaled8(a, b, n, r);
return;
}
#endif
int i = 0;
U q = quotient(r);
for (; i < n; ++i) a[i] = add(a[i], fixed_mul(b[i], r, q));
}
static void scale_copy(U *dst, const U *src, int length, U r) {
#if defined(__x86_64__) || defined(__i386__)
if (length >= 8 && vectorized()) {
scale_copy8(dst, src, length, r);
return;
}
#endif
U q = quotient(r);
int i = 0;
for (; i < length; ++i) dst[i] = fixed_mul(src[i], r, q);
}
void scale(U *a, int length, U r) const {
scale_copy(a, a, length, r);
}
struct CompositionDelete {
void operator()(unsigned int *p) const { ::operator delete[](p, std::align_val_t(64)); }
};
using CompositionWords = std::unique_ptr<unsigned int[], CompositionDelete>;
};
inline NTT &ntt_plan(int size) {
static NTT t(1, false);
if (static_cast<int>(t.roots.size()) < size) t = NTT(size, false);
return t;
}
}
inline void ntt(std::vector<mint> &a, bool inverse = false) {
assert(!a.empty() && !(a.size() & (a.size() - 1)) && a.size() <= (1U << 23));
int n = static_cast<int>(a.size());
auto &t = fps_detail::ntt_plan(std::max(1, n / 2));
auto reorder = [&] {
for (int i = 1, j = 0; i < n; ++i) {
int bit = n >> 1;
for (; j & bit; bit >>= 1) j ^= bit;
j ^= bit;
if (i < j) std::swap(a[i], a[j]);
}
};
static_assert(sizeof(mint) == sizeof(fps_detail::NTT::U) && std::is_trivially_copyable_v<mint>);
auto *words = reinterpret_cast<fps_detail::NTT::U *>(a.data());
if (inverse) {
reorder();
t.fft<true>(words, n);
t.scale(words, n, mint(n).inv().val());
} else {
t.fft<false>(words, n);
reorder();
}
}
#if defined(__GNUC__) && !defined(__clang__) && (defined(__x86_64__) || defined(__i386__))
#pragma GCC pop_options
#endif
#line 3 "test/ntt.test.cpp"
#include <bit>
#include <iostream>
#line 6 "test/ntt.test.cpp"
int main() {
std::cin.tie(nullptr)->sync_with_stdio(false);
int n, m;
std::cin >> n >> m;
int length = n + m - 1, size = std::bit_ceil(unsigned(length));
std::vector<fps_mint> a(size), b(size);
for (int i = 0; i < n; ++i) std::cin >> a[i];
for (int i = 0; i < m; ++i) std::cin >> b[i];
ntt(a);
ntt(b);
for (int i = 0; i < size; ++i) a[i] *= b[i];
ntt(a, true);
for (int i = 0; i < length; ++i) std::cout << a[i] << ' ';
std::cout << '\n';
}