cp-library
競技プログラミング用 C++ ライブラリ
test/convolution_light.test.cpp
test/convolution_light.test.cpp · GitHub で見る
依存ライブラリ
コード
#define PROBLEM "https://judge.yosupo.jp/problem/convolution_mod"
#include "../convolution/convolution_light.hpp"
#include <iostream>
#include <vector>
int main() {
std::cin.tie(nullptr)->sync_with_stdio(false);
int n, m;
std::cin >> n >> m;
std::vector<mint> a(n), b(m);
for (auto &x : a) std::cin >> x;
for (auto &x : b) std::cin >> x;
for (auto x : light_convolution::convolution(a, b)) std::cout << x << ' ';
std::cout << '\n';
}
#line 1 "test/convolution_light.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/convolution_mod"
#line 1 "convolution/convolution_light.hpp"
#line 1 "components/fps_light_ntt.hpp"
#include <algorithm>
#include <bit>
#include <cassert>
#include <cstddef>
#include <type_traits>
#include <vector>
#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 11 "components/fps_light_ntt.hpp"
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
#include <immintrin.h>
#endif
#if defined(__GNUC__) && !defined(__clang__) && (defined(__x86_64__) || defined(__i386__))
#pragma GCC push_options
#pragma GCC optimize("O3")
#endif
namespace fps_light_detail {
using std::size_t;
using Series = std::vector<mint>;
static_assert(sizeof(mint) == sizeof(unsigned) && std::is_trivially_copyable_v<mint>);
inline bool simd_available() {
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
return __builtin_cpu_supports("avx2");
#else
return false;
#endif
}
template<bool Inverse>
void scalar_stage(mint *a, size_t n, size_t h, const mint *roots) {
for (size_t start = 0; start < n; start += h * 2)
for (size_t j = 0; j < h; ++j) {
mint x = a[start + j], y = a[start + j + h];
if constexpr (Inverse) y *= roots[j];
a[start + j] = x + y;
a[start + j + h] = x - y;
if constexpr (!Inverse) a[start + j + h] *= roots[j];
}
}
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
__attribute__((target("avx2")))
inline __m256i product(__m256i x, __m256i y) {
const auto p = _mm256_set1_epi32(mint::mod());
auto even = _mm256_mul_epu32(x, y);
auto odd = _mm256_mul_epu32(_mm256_srli_epi64(x, 32), _mm256_srli_epi64(y, 32));
auto q = _mm256_mullo_epi32(_mm256_mullo_epi32(x, y), _mm256_set1_epi32(998244351));
even = _mm256_srli_epi64(_mm256_add_epi64(even, _mm256_mul_epu32(q, p)), 32);
odd = _mm256_add_epi64(odd, _mm256_mul_epu32(_mm256_srli_epi64(q, 32), p));
return _mm256_blend_epi32(even, odd, 0xaa);
}
__attribute__((target("avx2")))
inline __m256i fixed_product(__m256i x, __m256i w, __m256i q) {
auto even = _mm256_mul_epu32(x, q);
auto odd = _mm256_mul_epu32(_mm256_srli_epi64(x, 32), _mm256_srli_epi64(q, 32));
q = _mm256_blend_epi32(_mm256_srli_epi64(even, 32), odd, 0xaa);
return _mm256_sub_epi32(_mm256_mullo_epi32(x, w), _mm256_mullo_epi32(q, _mm256_set1_epi32(mint::mod())));
}
__attribute__((target("avx2")))
inline void pointwise_avx2(mint *a, const mint *b, size_t n) {
size_t i = 0;
for (; i + 8 <= n; i += 8) {
auto x = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + i));
auto y = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(b + i));
_mm256_storeu_si256(reinterpret_cast<__m256i *>(a + i), product(x, y));
}
for (; i < n; ++i) a[i] *= b[i];
}
__attribute__((target("avx2")))
inline void scale_avx2(mint *a, size_t n, mint factor) {
size_t i = 0;
unsigned value = factor.val();
const auto f = _mm256_set1_epi32(value);
const auto q = _mm256_set1_epi32((static_cast<unsigned long long>(value) << 32) / mint::mod());
for (; i + 8 <= n; i += 8) {
auto x = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + i));
_mm256_storeu_si256(reinterpret_cast<__m256i *>(a + i), fixed_product(x, f, q));
}
for (; i < n; ++i) a[i] *= factor;
}
__attribute__((target("avx2")))
inline __m256i sum_vec(__m256i x, __m256i y) {
auto z = _mm256_add_epi32(x, y);
return _mm256_min_epu32(z, _mm256_sub_epi32(z, _mm256_set1_epi32(2 * mint::mod())));
}
__attribute__((target("avx2")))
inline __m256i diff_vec(__m256i x, __m256i y) {
auto z = _mm256_sub_epi32(x, y);
return _mm256_min_epu32(z, _mm256_add_epi32(z, _mm256_set1_epi32(2 * mint::mod())));
}
template<int Half, bool Inverse>
__attribute__((target("avx2"), always_inline))
inline __m256i leaf_stage(__m256i x, const unsigned *w, const unsigned *q) {
constexpr int mask = Half == 4 ? 0xf0 : Half == 2 ? 0xcc : 0xaa;
__m256i factor{}, quotient{};
if constexpr (Half == 4) {
factor = _mm256_broadcastsi128_si256(_mm_loadu_si128(reinterpret_cast<const __m128i *>(w + 4)));
quotient = _mm256_broadcastsi128_si256(_mm_loadu_si128(reinterpret_cast<const __m128i *>(q + 4)));
}
if constexpr (Half == 2) {
factor = _mm256_broadcastq_epi64(_mm_loadl_epi64(reinterpret_cast<const __m128i *>(w + 2)));
quotient = _mm256_broadcastq_epi64(_mm_loadl_epi64(reinterpret_cast<const __m128i *>(q + 2)));
}
if constexpr (Inverse && Half > 1) x = _mm256_blend_epi32(x, fixed_product(x, factor, quotient), mask);
__m256i y;
if constexpr (Half == 4) y = _mm256_permute2x128_si256(x, x, 1);
if constexpr (Half == 2) y = _mm256_shuffle_epi32(x, 0x4e);
if constexpr (Half == 1) y = _mm256_shuffle_epi32(x, 0xb1);
auto high = diff_vec(y, x);
if constexpr (!Inverse && Half > 1) high = fixed_product(high, factor, quotient);
return _mm256_blend_epi32(sum_vec(x, y), high, mask);
}
template<bool Inverse>
__attribute__((target("avx2")))
void leaf8(mint *a, size_t n, const mint *roots) {
alignas(32) unsigned values[8], quotients[8];
for (int i = 0; i < 8; ++i) {
values[i] = roots[i].val();
quotients[i] = (static_cast<unsigned long long>(values[i]) << 32) / mint::mod();
}
for (size_t i = 0; i < n; i += 8) {
auto x = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + i));
x = leaf_stage<Inverse ? 1 : 4, Inverse>(x, values, quotients);
x = leaf_stage<2, Inverse>(x, values, quotients);
x = leaf_stage<Inverse ? 4 : 1, Inverse>(x, values, quotients);
_mm256_storeu_si256(reinterpret_cast<__m256i *>(a + i), x);
}
}
__attribute__((target("avx2")))
inline void expand_roots(mint *out, const mint *previous, size_t n, mint step) {
for (size_t i = 0; i < n; ++i) {
out[2 * i].v = previous[i].v;
out[2 * i + 1].v = mint::reduce(static_cast<unsigned long long>(previous[i].v) * step.v);
}
}
template<bool Inverse>
__attribute__((target("avx2")))
void vector_stage(mint *a, size_t n, size_t h, const mint *roots) {
for (size_t start = 0; start < n; start += h * 2)
for (size_t j = 0; j < h; j += 8) {
auto *first = reinterpret_cast<__m256i *>(a + start + j);
auto *last = reinterpret_cast<__m256i *>(a + start + j + h);
auto x = _mm256_loadu_si256(first), y = _mm256_loadu_si256(last);
auto w = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(roots + j));
if constexpr (Inverse) y = product(y, w);
auto sum = sum_vec(x, y), difference = diff_vec(x, y);
if constexpr (!Inverse) difference = product(difference, w);
_mm256_storeu_si256(first, sum);
_mm256_storeu_si256(last, difference);
}
}
#endif
struct Roots {
Series forward{0, 1}, inverse{0, 1};
void prepare(size_t n) {
size_t begin = forward.size();
if (begin >= n) return;
forward.resize(n), inverse.resize(n);
for (size_t h = begin; h < n; h *= 2) {
mint step = mint(3).pow((mint::mod() - 1) / (h * 2));
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
if (simd_available() && h >= 16) expand_roots(forward.data() + h, forward.data() + h / 2, h / 2, step);
else
#endif
for (size_t j = 0; j < h / 2; ++j) {
forward[h + 2 * j] = forward[h / 2 + j];
forward[h + 2 * j + 1] = forward[h / 2 + j] * step;
}
inverse[h] = 1;
for (size_t j = 1; j < h; ++j) inverse[h + j] = -forward[2 * h - j];
}
}
};
inline void pointwise(Series &a, const Series &b, [[maybe_unused]] bool vectorized = simd_available()) {
assert(a.size() <= b.size());
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
if (vectorized) return pointwise_avx2(a.data(), b.data(), a.size());
#endif
for (size_t i = 0; i < a.size(); ++i) a[i] *= b[i];
}
inline void scale(Series &a, mint factor, [[maybe_unused]] bool vectorized = simd_available()) {
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
if (vectorized) return scale_avx2(a.data(), a.size(), factor);
#endif
for (auto &x : a) x *= factor;
}
template<bool Inverse>
void stage(mint *a, size_t n, size_t h, const mint *roots, [[maybe_unused]] bool vectorized) {
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
if (vectorized && h >= 8) return vector_stage<Inverse>(a, n, h, roots);
#endif
scalar_stage<Inverse>(a, n, h, roots);
}
inline void transform(Series &a, bool inverse = false) {
size_t n = a.size();
assert(std::has_single_bit(n) && n <= (size_t(1) << 23));
static Roots roots;
roots.prepare(n);
bool vectorized = simd_available();
if (inverse) {
size_t h = 1;
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
if (vectorized && n >= 8) leaf8<true>(a.data(), n, roots.inverse.data()), h = 8;
#endif
for (; h < n; h *= 2) stage<true>(a.data(), n, h, roots.inverse.data() + h, vectorized);
scale(a, mint::raw(mint::mod() - (mint::mod() - 1) / n), vectorized);
} else {
size_t h = n / 2;
for (; h >= 8; h /= 2) stage<false>(a.data(), n, h, roots.forward.data() + h, vectorized);
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
if (vectorized && n >= 8) return leaf8<false>(a.data(), n, roots.forward.data());
#endif
for (; h; h /= 2) stage<false>(a.data(), n, h, roots.forward.data() + h, vectorized);
}
}
inline Series multiply(const Series &a, const Series &b) {
if (a.empty() || b.empty()) return {};
[[maybe_unused]] constexpr size_t limit = size_t(1) << 23;
assert(a.size() <= limit && b.size() <= limit && a.size() + b.size() - 1 <= limit);
size_t length = a.size() + b.size() - 1;
if (std::min(a.size(), b.size()) <= 32) {
Series result(length);
for (size_t i = 0; i < a.size(); ++i)
for (size_t j = 0; j < b.size(); ++j) result[i + j] += a[i] * b[j];
return result;
}
size_t n = std::bit_ceil(length);
Series x = a;
x.resize(n);
transform(x);
if (&a == &b) {
pointwise(x, x);
} else {
Series y = b;
y.resize(n);
transform(y);
pointwise(x, y);
}
transform(x, true);
x.resize(length);
return x;
}
}
#if defined(__GNUC__) && !defined(__clang__) && (defined(__x86_64__) || defined(__i386__))
#pragma GCC pop_options
#endif
#line 1 "convolution/convolution_fft.hpp"
#line 1 "components/fft.hpp"
#line 7 "components/fft.hpp"
#include <cmath>
#include <complex>
#line 10 "components/fft.hpp"
#include <numbers>
#line 12 "components/fft.hpp"
template <class Real>
void fft(std::vector<std::complex<Real>> &a, bool inverse = false) {
std::size_t n = a.size();
assert(!n || std::has_single_bit(n));
for (std::size_t i = 1, j = 0; i < n; ++i) {
std::size_t bit = n >> 1;
for (; j & bit; bit >>= 1) j ^= bit;
j ^= bit;
if (i < j) std::swap(a[i], a[j]);
}
for (std::size_t width = 2; width <= n; width *= 2) {
Real angle = (inverse ? -2 : 2) * std::numbers::pi_v<Real> / Real(width);
std::complex<Real> step(std::cos(angle), std::sin(angle));
for (std::size_t start = 0; start < n; start += width) {
std::complex<Real> w(1);
for (std::size_t j = 0; j < width / 2; ++j, w *= step) {
auto x = a[start + j], y = a[start + j + width / 2] * w;
a[start + j] = x + y;
a[start + j + width / 2] = x - y;
}
}
}
if (inverse)
for (auto &x : a) x /= Real(n);
}
#line 11 "convolution/convolution_fft.hpp"
template <class Real>
std::vector<std::complex<Real>> convolution_fft(const std::vector<std::complex<Real>> &a, const std::vector<std::complex<Real>> &b) {
if (a.empty() || b.empty()) return {};
std::size_t length = a.size() + b.size() - 1, n = std::bit_ceil(length);
std::vector<std::complex<Real>> x(a.begin(), a.end()), y(b.begin(), b.end());
x.resize(n);
y.resize(n);
fft(x, false);
if (&a == &b) y = x;
else fft(y, false);
for (std::size_t i = 0; i < n; ++i) x[i] *= y[i];
fft(x, true);
x.resize(length);
return x;
}
template <class Real>
requires std::is_floating_point_v<Real>
std::vector<Real> convolution_fft(const std::vector<Real> &a, const std::vector<Real> &b) {
std::vector<std::complex<Real>> x(a.begin(), a.end()), y(b.begin(), b.end());
auto values = &a == &b ? convolution_fft(x, x) : convolution_fft(x, y);
std::vector<Real> result;
result.reserve(values.size());
for (const auto &x : values) result.push_back(x.real());
return result;
}
#line 1 "convolution/convolution_mod.hpp"
#line 1 "components/ntt_mod.hpp"
#line 8 "components/ntt_mod.hpp"
#include <cstdint>
#line 10 "components/ntt_mod.hpp"
#if defined(__x86_64__) || defined(__i386__)
#include <immintrin.h>
#endif
namespace ntt_mod_detail {
template <std::uint32_t mod> struct Vector;
template <std::uint32_t mod>
std::uint32_t power(std::uint32_t a, std::uint32_t k) {
std::uint32_t result = 1;
for (; k; k >>= 1, a = std::uint64_t(a) * a % mod)
if (k & 1) result = std::uint64_t(result) * a % mod;
return result;
}
template <std::uint32_t mod>
void butterfly(std::uint32_t *a, std::size_t n, std::size_t half,
const std::uint32_t *roots, const std::uint32_t *quotients) {
for (std::size_t start = 0; start < n; start += 2 * half)
for (std::size_t j = 0; j < half; ++j) {
std::uint32_t x = a[start + j], v = a[start + j + half];
std::uint32_t q = std::uint64_t(v) * quotients[j] >> 32;
std::uint32_t y = v * roots[j] - q * mod;
if (y >= mod) y -= mod;
a[start + j] = x + y < mod ? x + y : x + y - mod;
a[start + j + half] = x >= y ? x - y : x + mod - y;
}
}
#if defined(__x86_64__) || defined(__i386__)
template <std::uint32_t mod>
__attribute__((target("avx2")))
void butterfly_avx2(std::uint32_t *a, std::size_t n, std::size_t half,
const std::uint32_t *roots, const std::uint32_t *quotients) {
using V = Vector<mod>;
for (std::size_t start = 0; start < n; start += 2 * half)
for (std::size_t j = 0; j < half; j += 8) {
auto x = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + start + j));
auto y = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + start + j + half));
auto r = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(roots + j));
auto q = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(quotients + j));
y = V::multiply(y, r, q);
_mm256_storeu_si256(reinterpret_cast<__m256i *>(a + start + j), V::add(x, y));
_mm256_storeu_si256(reinterpret_cast<__m256i *>(a + start + j + half), V::subtract(x, y));
}
}
#endif
template <std::uint32_t mod>
std::uint32_t shoup(std::uint32_t value, std::uint32_t factor, std::uint32_t quotient) {
std::uint32_t q = std::uint64_t(value) * quotient >> 32;
std::uint32_t result = value * factor - q * mod;
return result < mod ? result : result - mod;
}
template <std::uint32_t mod>
void forward_stage(std::uint32_t *a, std::size_t n, std::size_t half,
const std::uint32_t *roots, const std::uint32_t *quotients) {
for (std::size_t start = 0; start < n; start += 2 * half)
for (std::size_t j = 0; j < half; ++j) {
std::uint32_t x = a[start + j], y = a[start + j + half];
a[start + j] = x + y < mod ? x + y : x + y - mod;
std::uint32_t difference = x >= y ? x - y : x + mod - y;
a[start + j + half] = shoup<mod>(difference, roots[j], quotients[j]);
}
}
#if defined(__x86_64__) || defined(__i386__)
template <std::uint32_t mod>
__attribute__((target("avx2")))
void forward_stage_avx2(std::uint32_t *a, std::size_t n, std::size_t half,
const std::uint32_t *roots, const std::uint32_t *quotients) {
using V = Vector<mod>;
for (std::size_t start = 0; start < n; start += 2 * half)
for (std::size_t j = 0; j < half; j += 8) {
auto x = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + start + j));
auto y = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + start + j + half));
auto r = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(roots + j));
auto q = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(quotients + j));
auto difference = V::multiply(V::subtract(x, y), r, q);
_mm256_storeu_si256(reinterpret_cast<__m256i *>(a + start + j), V::add(x, y));
_mm256_storeu_si256(reinterpret_cast<__m256i *>(a + start + j + half), difference);
}
}
#endif
#if defined(__x86_64__) || defined(__i386__)
template <std::uint32_t mod>
struct Vector {
static constexpr std::uint32_t range = mod < (std::uint32_t(1) << 30) ? 2 * mod : mod;
__attribute__((target("avx2"), always_inline))
static inline __m256i add(__m256i x, __m256i y) {
auto sum = _mm256_add_epi32(x, y);
return _mm256_min_epu32(sum, _mm256_sub_epi32(sum, _mm256_set1_epi32(range)));
}
__attribute__((target("avx2"), always_inline))
static inline __m256i subtract(__m256i x, __m256i y) {
auto difference = _mm256_sub_epi32(x, y);
return _mm256_min_epu32(difference, _mm256_add_epi32(difference, _mm256_set1_epi32(range)));
}
__attribute__((target("avx2"), always_inline))
static inline __m256i multiply(__m256i x, __m256i factor, __m256i quotient) {
auto even = _mm256_mul_epu32(x, quotient);
auto odd = _mm256_mul_epu32(_mm256_srli_epi64(x, 32), _mm256_srli_epi64(quotient, 32));
auto q = _mm256_blend_epi32(_mm256_srli_epi64(even, 32), odd, 0xaa);
auto result = _mm256_sub_epi32(_mm256_mullo_epi32(x, factor), _mm256_mullo_epi32(q, _mm256_set1_epi32(mod)));
if constexpr (range == 2 * mod) return result;
return _mm256_min_epu32(result, _mm256_sub_epi32(result, _mm256_set1_epi32(mod)));
}
};
template <std::uint32_t mod, bool inverse>
__attribute__((target("avx2")))
void leaf_avx2(std::uint32_t *a, std::size_t n, const std::uint32_t *roots, const std::uint32_t *quotients) {
using V = Vector<mod>;
auto r4 = _mm256_broadcastsi128_si256(_mm_loadu_si128(reinterpret_cast<const __m128i *>(roots + 4)));
auto q4 = _mm256_broadcastsi128_si256(_mm_loadu_si128(reinterpret_cast<const __m128i *>(quotients + 4)));
auto r2 = _mm256_setr_epi32(roots[2], roots[3], roots[2], roots[3], roots[2], roots[3], roots[2], roots[3]);
auto q2 = _mm256_setr_epi32(quotients[2], quotients[3], quotients[2], quotients[3], quotients[2], quotients[3], quotients[2], quotients[3]);
for (std::size_t i = 0; i < n; i += 8) {
auto x = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + i));
if constexpr (!inverse) {
auto y = _mm256_permute2x128_si256(x, x, 1);
x = _mm256_blend_epi32(V::add(x, y), V::multiply(V::subtract(y, x), r4, q4), 0xf0);
y = _mm256_shuffle_epi32(x, 0x4e);
x = _mm256_blend_epi32(V::add(x, y), V::multiply(V::subtract(y, x), r2, q2), 0xcc);
y = _mm256_shuffle_epi32(x, 0xb1);
x = _mm256_blend_epi32(V::add(x, y), V::subtract(y, x), 0xaa);
} else {
auto y = _mm256_shuffle_epi32(x, 0xb1);
x = _mm256_blend_epi32(V::add(x, y), V::subtract(y, x), 0xaa);
x = _mm256_blend_epi32(x, V::multiply(x, r2, q2), 0xcc);
y = _mm256_shuffle_epi32(x, 0x4e);
x = _mm256_blend_epi32(V::add(x, y), V::subtract(y, x), 0xcc);
x = _mm256_blend_epi32(x, V::multiply(x, r4, q4), 0xf0);
y = _mm256_permute2x128_si256(x, x, 1);
x = _mm256_blend_epi32(V::add(x, y), V::subtract(y, x), 0xf0);
}
_mm256_storeu_si256(reinterpret_cast<__m256i *>(a + i), x);
}
}
template <std::uint32_t mod, bool inverse>
__attribute__((target("avx2")))
void pair_avx2(std::uint32_t *a, std::size_t n, std::size_t h, const std::uint32_t *roots, const std::uint32_t *quotients) {
using V = Vector<mod>;
for (std::size_t start = 0; start < n; start += 4 * h)
for (std::size_t j = 0; j < h; j += 8) {
auto x0 = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + start + j));
auto x1 = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + start + j + h));
auto x2 = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + start + j + 2 * h));
auto x3 = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + start + j + 3 * h));
auto r0 = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(roots + 2 * h + j));
auto q0 = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(quotients + 2 * h + j));
auto r1 = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(roots + 3 * h + j));
auto q1 = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(quotients + 3 * h + j));
auto r2 = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(roots + h + j));
auto q2 = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(quotients + h + j));
if constexpr (!inverse) {
auto y0 = V::add(x0, x2), y1 = V::add(x1, x3);
auto y2 = V::multiply(V::subtract(x0, x2), r0, q0);
auto y3 = V::multiply(V::subtract(x1, x3), r1, q1);
x0 = V::add(y0, y1); x1 = V::multiply(V::subtract(y0, y1), r2, q2);
x2 = V::add(y2, y3); x3 = V::multiply(V::subtract(y2, y3), r2, q2);
} else {
x1 = V::multiply(x1, r2, q2); x3 = V::multiply(x3, r2, q2);
auto y0 = V::add(x0, x1), y1 = V::subtract(x0, x1);
auto y2 = V::multiply(V::add(x2, x3), r0, q0);
auto y3 = V::multiply(V::subtract(x2, x3), r1, q1);
x0 = V::add(y0, y2); x2 = V::subtract(y0, y2);
x1 = V::add(y1, y3); x3 = V::subtract(y1, y3);
}
_mm256_storeu_si256(reinterpret_cast<__m256i *>(a + start + j), x0);
_mm256_storeu_si256(reinterpret_cast<__m256i *>(a + start + j + h), x1);
_mm256_storeu_si256(reinterpret_cast<__m256i *>(a + start + j + 2 * h), x2);
_mm256_storeu_si256(reinterpret_cast<__m256i *>(a + start + j + 3 * h), x3);
}
}
#endif
#if defined(__x86_64__) || defined(__i386__)
template <std::uint32_t mod>
__attribute__((target("avx2")))
void build_roots_avx2(std::uint32_t *roots, std::uint32_t *quotients, std::size_t half, std::uint32_t step) {
std::uint32_t initial[8], w = 1;
for (int i = 0; i < 8; ++i) {
initial[i] = w;
w = std::uint64_t(w) * step % mod;
}
auto x = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(initial));
auto factor = _mm256_set1_epi32(w), quotient = _mm256_set1_epi32((std::uint64_t(w) << 32) / mod);
for (std::size_t j = 0; j < half; j += 8) {
x = _mm256_min_epu32(x, _mm256_sub_epi32(x, _mm256_set1_epi32(mod)));
_mm256_storeu_si256(reinterpret_cast<__m256i *>(roots + j), x);
for (int k = 0; k < 8; ++k) quotients[j + k] = (std::uint64_t(roots[j + k]) << 32) / mod;
x = Vector<mod>::multiply(x, factor, quotient);
}
}
#endif
template <std::uint32_t mod, std::uint32_t root>
struct Plan {
std::vector<std::uint32_t> roots, quotients;
explicit Plan(std::size_t n, [[maybe_unused]] bool vectorized = false) : roots(n), quotients(n) {
static_assert(mod > 2 && mod < (std::uint32_t(1) << 31));
assert(std::has_single_bit(n) && (mod - 1) % n == 0);
if (n == 1) return;
std::size_t half = n / 2;
std::uint32_t step = power<mod>(root, (mod - 1) / n), w = 1;
#if defined(__x86_64__) || defined(__i386__)
if (vectorized && half >= 8) build_roots_avx2<mod>(roots.data() + half, quotients.data() + half, half, step);
else
#endif
for (std::size_t j = 0; j < half; ++j) {
roots[half + j] = w;
quotients[half + j] = (std::uint64_t(w) << 32) / mod;
w = std::uint64_t(w) * step % mod;
}
for (half /= 2; half; half /= 2)
for (std::size_t j = 0; j < half; ++j) {
roots[half + j] = roots[2 * half + 2 * j];
quotients[half + j] = quotients[2 * half + 2 * j];
}
}
};
template <std::uint32_t mod, std::uint32_t root>
void forward(std::vector<std::uint32_t> &a, const Plan<mod, root> &plan, [[maybe_unused]] bool vectorized) {
std::size_t half = a.size() / 2;
#if defined(__x86_64__) || defined(__i386__)
if (vectorized && a.size() >= 8) {
for (; half >= 16; half /= 4) pair_avx2<mod, false>(a.data(), a.size(), half / 2, plan.roots.data(), plan.quotients.data());
if (half == 8) forward_stage_avx2<mod>(a.data(), a.size(), half, plan.roots.data() + half, plan.quotients.data() + half);
leaf_avx2<mod, false>(a.data(), a.size(), plan.roots.data(), plan.quotients.data());
return;
}
#endif
for (; half; half /= 2) forward_stage<mod>(a.data(), a.size(), half, plan.roots.data() + half, plan.quotients.data() + half);
}
template <std::uint32_t mod, std::uint32_t root>
void backward(std::vector<std::uint32_t> &a, const Plan<mod, root> &plan, [[maybe_unused]] bool vectorized) {
std::size_t half = 1;
#if defined(__x86_64__) || defined(__i386__)
if (vectorized && a.size() >= 8) {
leaf_avx2<mod, true>(a.data(), a.size(), plan.roots.data(), plan.quotients.data());
half = 8;
for (; 2 * half < a.size(); half *= 4) pair_avx2<mod, true>(a.data(), a.size(), half, plan.roots.data(), plan.quotients.data());
if (half < a.size()) butterfly_avx2<mod>(a.data(), a.size(), half, plan.roots.data() + half, plan.quotients.data() + half);
return;
}
#endif
for (; half < a.size(); half *= 2) butterfly<mod>(a.data(), a.size(), half, plan.roots.data() + half, plan.quotients.data() + half);
}
#if defined(__x86_64__) || defined(__i386__)
template <std::uint32_t mod>
__attribute__((target("avx2")))
void scale_avx2(std::uint32_t *a, std::size_t n, std::uint32_t factor, std::uint32_t quotient) {
auto r = _mm256_set1_epi32(factor), q = _mm256_set1_epi32(quotient);
for (std::size_t i = 0; i < n; i += 8) {
auto x = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + i));
auto y = Vector<mod>::multiply(x, r, q);
_mm256_storeu_si256(reinterpret_cast<__m256i *>(a + i), _mm256_min_epu32(y, _mm256_sub_epi32(y, _mm256_set1_epi32(mod))));
}
}
#endif
template <std::uint32_t mod>
void scale(std::vector<std::uint32_t> &a, std::uint32_t factor, [[maybe_unused]] bool vectorized) {
std::uint32_t quotient = (std::uint64_t(factor) << 32) / mod;
#if defined(__x86_64__) || defined(__i386__)
if (vectorized && a.size() >= 8) {
scale_avx2<mod>(a.data(), a.size(), factor, quotient);
return;
}
#endif
for (auto &x : a) x = shoup<mod>(x, factor, quotient);
}
inline void bit_reverse(std::vector<std::uint32_t> &a) {
for (std::size_t i = 1, j = 0; i < a.size(); ++i) {
std::size_t bit = a.size() >> 1;
for (; j & bit; bit >>= 1) j ^= bit;
j ^= bit;
if (i < j) std::swap(a[i], a[j]);
}
}
template <std::uint32_t mod, std::uint32_t root>
void transform(std::vector<std::uint32_t> &a, bool inverse, bool vectorized) {
Plan<mod, root> plan(a.size(), vectorized);
if (inverse) {
bit_reverse(a);
backward(a, plan, vectorized);
std::reverse(a.begin() + 1, a.end());
scale<mod>(a, mod - (mod - 1) / a.size(), vectorized);
} else {
forward(a, plan, vectorized);
if constexpr (mod < (std::uint32_t(1) << 30)) for (auto &x : a) if (x >= mod) x -= mod;
bit_reverse(a);
}
}
template <std::uint32_t mod>
constexpr std::uint32_t negative_inverse = [] {
std::uint32_t value = mod;
for (int i = 0; i < 5; ++i) value *= 2 - mod * value;
return 0 - value;
}();
#if defined(__x86_64__) || defined(__i386__)
template <std::uint32_t mod>
__attribute__((target("avx2")))
void pointwise_avx2(std::uint32_t *a, const std::uint32_t *b, std::size_t n) {
auto p = _mm256_set1_epi32(mod), inverse = _mm256_set1_epi32(negative_inverse<mod>);
for (std::size_t i = 0; i < n; i += 8) {
auto x = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + i));
auto y = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(b + i));
auto even = _mm256_mul_epu32(x, y);
auto odd = _mm256_mul_epu32(_mm256_srli_epi64(x, 32), _mm256_srli_epi64(y, 32));
auto qe = _mm256_mul_epu32(even, inverse), qo = _mm256_mul_epu32(odd, inverse);
even = _mm256_srli_epi64(_mm256_add_epi64(even, _mm256_mul_epu32(qe, p)), 32);
odd = _mm256_add_epi64(odd, _mm256_mul_epu32(qo, p));
x = _mm256_blend_epi32(even, odd, 0xaa);
x = _mm256_min_epu32(x, _mm256_sub_epi32(x, p));
_mm256_storeu_si256(reinterpret_cast<__m256i *>(a + i), x);
}
}
#endif
template <std::uint32_t mod>
void pointwise(std::vector<std::uint32_t> &a, const std::vector<std::uint32_t> &b, [[maybe_unused]] bool vectorized) {
#if defined(__x86_64__) || defined(__i386__)
if (vectorized && a.size() >= 8) {
pointwise_avx2<mod>(a.data(), b.data(), a.size());
return;
}
#endif
for (std::size_t i = 0; i < a.size(); ++i) {
std::uint64_t product = std::uint64_t(a[i]) * b[i];
std::uint32_t q = std::uint32_t(product) * negative_inverse<mod>;
std::uint32_t value = (product + std::uint64_t(q) * mod) >> 32;
a[i] = value < mod ? value : value - mod;
}
}
}
template <std::uint32_t mod, std::uint32_t root>
void ntt_mod(std::vector<std::uint32_t> &a, bool inverse = false) {
bool vectorized = false;
#if defined(__x86_64__) || defined(__i386__)
vectorized = __builtin_cpu_supports("avx2");
#endif
ntt_mod_detail::transform<mod, root>(a, inverse, vectorized);
}
#line 9 "convolution/convolution_mod.hpp"
#include <limits>
#include <utility>
#line 12 "convolution/convolution_mod.hpp"
namespace convolution_mod_detail {
using Word = std::uint32_t;
template <Word mod, Word root, class T>
std::vector<Word> multiply(const std::vector<T> &a, const std::vector<T> &b, Word factor = 1) {
const auto &small = a.size() <= b.size() ? a : b;
const auto &large = a.size() <= b.size() ? b : a;
std::size_t length = a.size() + b.size() - 1;
bool blocked = small.size() > 32 && large.size() / small.size() >= 8, square = &a == &b;
std::size_t n = std::bit_ceil(blocked ? 2 * small.size() - 1 : length), block = n - small.size() + 1;
std::vector<Word> x(n), y(square ? 0 : n);
auto reduce = [](T value) {
__int128_t r = __int128_t(value) % mod;
return Word(r < 0 ? r + mod : r);
};
for (std::size_t i = 0; i < small.size(); ++i) x[i] = reduce(small[i]);
if (small.size() <= 32) {
if (!square) for (std::size_t i = 0; i < large.size(); ++i) y[i] = reduce(large[i]);
const auto &other = square ? x : y;
std::vector<Word> result(length);
for (std::size_t i = 0; i < small.size(); ++i)
for (std::size_t j = 0; j < large.size(); ++j) result[i + j] = (result[i + j] + std::uint64_t(x[i]) * other[j]) % mod;
if (factor != 1) for (auto &v : result) v = std::uint64_t(v) * factor % mod;
return result;
}
bool vectorized = false;
#if defined(__x86_64__) || defined(__i386__)
vectorized = __builtin_cpu_supports("avx2");
#endif
ntt_mod_detail::Plan<mod, root> plan(n, vectorized);
ntt_mod_detail::forward(x, plan, vectorized);
Word scale = (std::uint64_t(1) << 32) % mod * (mod - (mod - 1) / n) % mod * factor % mod;
std::vector<Word> result(blocked ? length : 0);
auto &work = square ? x : y;
for (std::size_t start = 0; start < large.size(); start += block) {
std::size_t size = std::min(block, large.size() - start);
if (!square) {
for (std::size_t i = 0; i < size; ++i) y[i] = reduce(large[start + i]);
std::fill(y.begin() + size, y.end(), 0);
ntt_mod_detail::forward(y, plan, vectorized);
}
ntt_mod_detail::pointwise<mod>(work, x, vectorized);
ntt_mod_detail::backward(work, plan, vectorized);
std::reverse(work.begin() + 1, work.end());
ntt_mod_detail::scale<mod>(work, scale, vectorized);
if (!blocked) {
work.resize(length);
return std::move(work);
}
for (std::size_t i = 0; i < size + small.size() - 1; ++i) {
Word sum = result[start + i] + work[i];
result[start + i] = sum < mod ? sum : sum - mod;
}
}
return result;
}
inline constexpr Word p = 754974721, q = 1224736769, r = 2013265921;
inline constexpr __uint128_t modulus = __uint128_t(p) * q * r;
struct Reconstruction { std::uint64_t low; Word high; };
template <Word P = p, Word Q = q, Word R = r>
inline Reconstruction reconstruct(Word a, Word b, Word c) {
static const Word inverse_p = ntt_mod_detail::power<Q>(P, Q - 2);
static const Word inverse_pq = ntt_mod_detail::power<R>(std::uint64_t(P) * Q % R, R - 2);
std::uint64_t x = std::uint64_t(b + Q - a) * inverse_p % Q;
std::uint64_t low = a + std::uint64_t(P) * x;
std::uint64_t y = std::uint64_t(c + R - low % R) * inverse_pq % R;
return {low, Word(y)};
}
inline __uint128_t restore(Word a, Word b, Word c) {
auto [low, high] = reconstruct(a, b, c);
return low + __uint128_t(p) * q * high;
}
#if defined(__x86_64__) || defined(__i386__)
template <Word P, Word Q, Word R>
__attribute__((target("avx2")))
void combine(Word *a, const Word *b, const Word *c, std::size_t n) {
static_assert(P > 256 && Q > 256 && R > 256 && P + std::uint64_t(Q) + R < (std::uint64_t(1) << 32));
constexpr Word mod = 1000000007;
constexpr auto form = [](std::uint64_t x) { return Word(x % mod * ((std::uint64_t(1) << 32) % mod) % mod); };
constexpr Word place[3] = {form(std::uint64_t(Q) * R), form(std::uint64_t(P) * R), form(std::uint64_t(P) * Q)};
constexpr Word wrap = form(mod - std::uint64_t(P) * Q % mod * R % mod);
const auto p = _mm256_set1_epi32(mod), inverse = _mm256_set1_epi32(ntt_mod_detail::negative_inverse<mod>);
std::size_t i = 0;
for (; n - i >= 8; i += 8) {
auto x = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + i));
auto y = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(b + i));
auto z = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(c + i));
auto quotient = [&](__m256i a, __m256i b, __m256i c) __attribute__((target("avx2"), always_inline)) {
auto s = _mm256_mul_epu32(a, _mm256_set1_epi32((std::uint64_t(1) << 40) / P + 1));
s = _mm256_add_epi64(s, _mm256_mul_epu32(b, _mm256_set1_epi32((std::uint64_t(1) << 40) / Q + 1)));
s = _mm256_add_epi64(s, _mm256_mul_epu32(c, _mm256_set1_epi32((std::uint64_t(1) << 40) / R + 1)));
return _mm256_srli_epi64(_mm256_add_epi64(s, _mm256_set1_epi64x(std::uint64_t(1) << 39)), 40);
};
auto reduce = [&](__m256i a, __m256i b, __m256i c, __m256i t) __attribute__((target("avx2"), always_inline)) {
auto s = _mm256_add_epi64(_mm256_mul_epu32(a, _mm256_set1_epi32(place[0])), _mm256_mul_epu32(b, _mm256_set1_epi32(place[1])));
s = _mm256_add_epi64(s, _mm256_mul_epu32(c, _mm256_set1_epi32(place[2])));
s = _mm256_add_epi64(s, _mm256_mul_epu32(t, _mm256_set1_epi32(wrap)));
return _mm256_add_epi64(s, _mm256_mul_epu32(_mm256_mul_epu32(s, inverse), p));
};
auto xo = _mm256_srli_epi64(x, 32), yo = _mm256_srli_epi64(y, 32), zo = _mm256_srli_epi64(z, 32);
auto even = reduce(x, y, z, quotient(x, y, z)), odd = reduce(xo, yo, zo, quotient(xo, yo, zo));
x = _mm256_blend_epi32(_mm256_srli_epi64(even, 32), odd, 0xaa);
x = _mm256_min_epu32(x, _mm256_sub_epi32(x, p));
_mm256_storeu_si256(reinterpret_cast<__m256i *>(a + i), x);
}
for (; i < n; ++i) {
std::uint64_t t = (std::uint64_t(a[i]) * ((std::uint64_t(1) << 40) / P + 1) + std::uint64_t(b[i]) * ((std::uint64_t(1) << 40) / Q + 1) + std::uint64_t(c[i]) * ((std::uint64_t(1) << 40) / R + 1) + (std::uint64_t(1) << 39)) >> 40;
std::uint64_t sum = std::uint64_t(a[i]) * place[0] + std::uint64_t(b[i]) * place[1] + std::uint64_t(c[i]) * place[2] + std::uint64_t(t) * wrap;
auto q = Word(sum) * ntt_mod_detail::negative_inverse<mod>;
Word value = (sum + std::uint64_t(q) * mod) >> 32;
a[i] = value < mod ? value : value - mod;
}
}
#endif
template <Word P, Word PR, Word Q, Word QR, Word R, Word RR>
std::vector<Word> modular(const std::vector<Word> &a, const std::vector<Word> &b, Word mod) {
#if defined(__x86_64__) || defined(__i386__)
if constexpr (P + std::uint64_t(Q) + R < (std::uint64_t(1) << 32)) {
if (mod == 1000000007 && __uint128_t(std::min(a.size(), b.size())) * (mod - 1) * (mod - 1) * 4 < __uint128_t(P) * Q * R && __builtin_cpu_supports("avx2")) {
auto x = multiply<P, PR>(a, b, ntt_mod_detail::power<P>(std::uint64_t(Q) * R % P, P - 2));
auto y = multiply<Q, QR>(a, b, ntt_mod_detail::power<Q>(std::uint64_t(P) * R % Q, Q - 2));
auto z = multiply<R, RR>(a, b, ntt_mod_detail::power<R>(std::uint64_t(P) * Q % R, R - 2));
combine<P, Q, R>(x.data(), y.data(), z.data(), x.size());
return x;
}
}
#endif
auto x = multiply<P, PR>(a, b);
auto y = multiply<Q, QR>(a, b);
auto z = multiply<R, RR>(a, b);
std::uint64_t pq_mod = std::uint64_t(P) * Q % mod;
for (std::size_t i = 0; i < x.size(); ++i) {
auto [low, high] = reconstruct<P, Q, R>(x[i], y[i], z[i]);
x[i] = (low % mod + pq_mod * high) % mod;
}
return x;
}
}
template <class T>
std::vector<T> convolution_integer(const std::vector<T> &a, const std::vector<T> &b) {
static_assert(std::numeric_limits<T>::is_integer && sizeof(T) <= 8);
if (a.empty() || b.empty()) return {};
assert(a.size() + b.size() - 1 <= std::size_t(1) << 24);
auto x = convolution_mod_detail::multiply<convolution_mod_detail::p, 11>(a, b);
auto y = convolution_mod_detail::multiply<convolution_mod_detail::q, 3>(a, b);
auto z = convolution_mod_detail::multiply<convolution_mod_detail::r, 31>(a, b);
std::vector<T> result(x.size());
for (std::size_t i = 0; i < x.size(); ++i) {
auto value = convolution_mod_detail::restore(x[i], y[i], z[i]);
if constexpr (std::numeric_limits<T>::is_signed) {
__int128_t v = value > convolution_mod_detail::modulus / 2 ? __int128_t(value) - __int128_t(convolution_mod_detail::modulus) : __int128_t(value);
assert(std::numeric_limits<T>::lowest() <= v && v <= std::numeric_limits<T>::max());
result[i] = T(v);
} else {
assert(value <= std::numeric_limits<T>::max());
result[i] = T(value);
}
}
return result;
}
template <class T>
std::vector<std::uint32_t> convolution_mod(const std::vector<T> &a, const std::vector<T> &b, std::uint32_t mod) {
static_assert(std::numeric_limits<T>::is_integer && sizeof(T) <= 8);
assert(mod > 0);
if (a.empty() || b.empty()) return {};
assert(a.size() + b.size() - 1 <= std::size_t(1) << 24);
if (mod == 1) return std::vector<std::uint32_t>(a.size() + b.size() - 1);
std::vector<std::uint32_t> x(a.size()), y(&a == &b ? 0 : b.size());
auto reduce = [=](T value) {
__int128_t r = __int128_t(value) % mod;
return std::uint32_t(r < 0 ? r + mod : r);
};
std::transform(a.begin(), a.end(), x.begin(), reduce);
if (&a != &b) std::transform(b.begin(), b.end(), y.begin(), reduce);
const auto &other = &a == &b ? x : y;
if (std::min(x.size(), other.size()) <= 32) {
std::vector<std::uint32_t> result(x.size() + other.size() - 1);
for (std::size_t i = 0; i < x.size(); ++i)
for (std::size_t j = 0; j < other.size(); ++j) result[i + j] = (result[i + j] + std::uint64_t(x[i]) * other[j]) % mod;
return result;
}
if (mod == 998244353 && x.size() + other.size() - 1 <= std::size_t(1) << 23) return convolution_mod_detail::multiply<998244353, 3>(x, other);
constexpr std::uint32_t P = 167772161, Q = 469762049, R = 754974721;
if (__uint128_t(std::min(x.size(), other.size())) * (mod - 1) * (mod - 1) < __uint128_t(P) * Q * R)
return convolution_mod_detail::modular<P, 3, Q, 3, R, 11>(x, other, mod);
return convolution_mod_detail::modular<convolution_mod_detail::p, 11, convolution_mod_detail::q, 3, convolution_mod_detail::r, 31>(x, other, mod);
}
#line 1 "convolution/convolution_ring.hpp"
#line 6 "convolution/convolution_ring.hpp"
#include <span>
#line 8 "convolution/convolution_ring.hpp"
namespace convolution_ring_detail {
template <class T>
std::vector<T> multiply(std::span<const T> a, std::span<const T> b) {
if (a.empty() || b.empty()) return {};
std::vector<T> result(a.size() + b.size() - 1, T(0));
if (std::min(a.size(), b.size()) <= 32) {
for (std::size_t i = 0; i < a.size(); ++i)
for (std::size_t j = 0; j < b.size(); ++j) result[i + j] += a[i] * b[j];
return result;
}
std::size_t half = std::max(a.size(), b.size()) / 2;
auto al = a.first(std::min(a.size(), half)), ar = a.subspan(al.size());
auto bl = b.first(std::min(b.size(), half)), br = b.subspan(bl.size());
auto low = multiply(al, bl), high = multiply(ar, br);
std::vector<T> x(std::max(al.size(), ar.size()), T(0)), y(std::max(bl.size(), br.size()), T(0));
for (std::size_t i = 0; i < al.size(); ++i) x[i] += al[i];
for (std::size_t i = 0; i < ar.size(); ++i) x[i] += ar[i];
for (std::size_t i = 0; i < bl.size(); ++i) y[i] += bl[i];
for (std::size_t i = 0; i < br.size(); ++i) y[i] += br[i];
auto middle = multiply<T>(x, y);
for (std::size_t i = 0; i < low.size(); ++i) result[i] += low[i], middle[i] -= low[i];
for (std::size_t i = 0; i < high.size(); ++i) result[i + 2 * half] += high[i], middle[i] -= high[i];
for (std::size_t i = 0; i < middle.size() && i + half < result.size(); ++i) result[i + half] += middle[i];
return result;
}
}
template <class T>
std::vector<T> convolution_ring(const std::vector<T> &a, const std::vector<T> &b) {
return convolution_ring_detail::multiply<T>(a, b);
}
#line 12 "convolution/convolution_light.hpp"
namespace light_convolution {
namespace detail {
template<class T> struct IsComplex : std::false_type {};
template<class T> struct IsComplex<std::complex<T>> : std::true_type {};
}
template<class T>
std::vector<T> convolution(const std::vector<T> &a, const std::vector<T> &b) {
if constexpr (requires { T::mod(); T::raw(0); a[0].val(); }) {
if constexpr (std::is_same_v<T, mint>) return fps_light_detail::multiply(a, b);
using M = std::conditional_t<T::mod() == 998244353, mint, std::uint32_t>;
auto convert = [](const std::vector<T> &input) {
std::vector<M> result;
result.reserve(input.size());
for (auto x : input) {
if constexpr (std::is_same_v<M, mint>) result.push_back(mint::raw(x.val()));
else result.push_back(x.val());
}
return result;
};
auto x = convert(a), y = &a == &b ? std::vector<M>{} : convert(b);
const auto &other = &a == &b ? x : y;
auto values = [&] {
if constexpr (std::is_same_v<M, mint>) return fps_light_detail::multiply(x, other);
else return convolution_mod(x, other, T::mod());
}();
std::vector<T> result;
result.reserve(values.size());
for (auto value : values) {
if constexpr (std::is_same_v<M, mint>) result.push_back(T::raw(value.val()));
else result.push_back(T::raw(value));
}
return result;
} else if constexpr (std::numeric_limits<T>::is_integer && sizeof(T) <= 8) return convolution_integer(a, b);
else if constexpr (std::is_floating_point_v<T> || detail::IsComplex<T>::value) return convolution_fft(a, b);
else return convolution_ring(a, b);
}
}
#line 3 "test/convolution_light.test.cpp"
#include <iostream>
#line 5 "test/convolution_light.test.cpp"
int main() {
std::cin.tie(nullptr)->sync_with_stdio(false);
int n, m;
std::cin >> n >> m;
std::vector<mint> a(n), b(m);
for (auto &x : a) std::cin >> x;
for (auto &x : b) std::cin >> x;
for (auto x : light_convolution::convolution(a, b)) std::cout << x << ' ';
std::cout << '\n';
}