test/math/ntt.test.cpp

• View this file on GitHub
• Last update: 2025-10-26 10:54:48+08:00
• Problem: https://judge.yosupo.jp/problem/convolution_mod

Depends on

• src/math/modint.hpp
• src/math/ntt.hpp
• src/misc/read.hpp

Code

// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/convolution_mod
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using ull = unsigned long long;

#include "../../src/math/ntt.hpp"
#include "../../src/misc/read.hpp"

int main() {
  cin.tie(nullptr)->sync_with_stdio(false);
  int n = read(), m = read();
  Poly F, G;
  F.data.resize(n);
  for (int i = 0; i < n; i++) F.data[i] = read();
  G.data.resize(m);
  for (int i = 0; i < m; i++) G.data[i] = read();
  Poly H = F * G;
  for (int i = 0; i < n + m - 1; i++) {
    cout << H.data[i].data << " \n"[i == n + m - 2];
  }
  return 0;
}

#line 1 "test/math/ntt.test.cpp"
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/convolution_mod
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using ull = unsigned long long;

#line 1 "src/math/modint.hpp"
template <unsigned MOD>
struct ModInt {
  unsigned data;
  ModInt(ll v = 0) : data(norm(v % MOD)) {}
  ModInt operator-() const { return MOD - data; }
  ModInt& operator+=(ModInt rhs) { return data = norm(data + rhs.data), *this; }
  ModInt& operator-=(ModInt rhs) { return data = norm(data - rhs.data), *this; }
  ModInt& operator*=(ModInt rhs) {
    return data = ull(data) * rhs.data % MOD, *this;
  }
  ModInt& operator/=(ModInt rhs) {
    return data = ull(data) * rhs.inv() % MOD, *this;
  }
  friend ModInt operator+(ModInt lhs, ModInt rhs) { return lhs += rhs; }
  friend ModInt operator-(ModInt lhs, ModInt rhs) { return lhs -= rhs; }
  friend ModInt operator*(ModInt lhs, ModInt rhs) { return lhs *= rhs; }
  friend ModInt operator/(ModInt lhs, ModInt rhs) { return lhs /= rhs; }
  unsigned inv() const {
    ll x, y;  // Inverse does not exist if gcd(data, MOD) != 1.
    assert(exgcd(data, MOD, x, y) == 1);
    return norm(x);
  }
  ModInt pow(ull n) const { return pow_mod(data, n, MOD); }
  static ll exgcd(ll a, ll b, ll& x, ll& y) {
    x = 1, y = 0;
    ll x1 = 0, y1 = 1;
    while (b) {
      ll q = a / b;
      swap(a -= q * b, b);
      swap(x -= q * x1, x1);
      swap(y -= q * y1, y1);
    }
    return a;
  }
  static unsigned pow_mod(unsigned a, ull n, unsigned p) {
    unsigned ret = 1;
    for (; n; n /= 2) {
      if (n & 1) ret = ull(ret) * a % p;
      a = ull(a) * a % p;
    }
    return ret;
  }

 private:
  static unsigned norm(unsigned x) {
    if ((x >> (8 * sizeof(unsigned) - 1)) & 1) x += MOD;
    return x >= MOD ? x -= MOD : x;
  }
};
constexpr unsigned MOD = 998244353;
using mint = ModInt<MOD>;
#line 2 "src/math/ntt.hpp"
template <int MOD = 998244353, int G = 114514, int GI = 137043501>
struct Poly {
  vector<mint> data;
  friend Poly operator*(const Poly& a, const Poly& b) {
    int n = 1, new_sz = a.data.size() + b.data.size() - 1;
    while (n < int(a.data.size() + b.data.size())) n *= 2;
    vector<mint> A = a.data, B = b.data;
    A.resize(n), B.resize(n);
    ntt(A, 1), ntt(B, 1);
    for (int i = 0; i < n; i++) A[i] *= B[i];
    ntt(A, -1), A.resize(new_sz);
    return {A};
  }

 private:
  static void ntt(vector<mint>& a, int type) {
    int len = a.size();
    vector<int> rev(len);
    for (int i = 0; i < len; i++) {
      rev[i] = rev[i >> 1] >> 1;
      if (i & 1) rev[i] |= len >> 1;
    }
    for (int i = 0; i < len; i++)
      if (i < rev[i]) swap(a[i], a[rev[i]]);
    for (int n = 2, m = 1; n <= len; n *= 2, m *= 2) {
      mint step = mint(type == 1 ? G : GI).pow((MOD - 1) / n);
      for (int j = 0; j < len; j += n) {
        mint w = 1;
        for (int k = j; k < j + m; k++) {
          mint u = a[k], v = w * a[k + m];
          a[k] = u + v, a[k + m] = u - v, w *= step;
        }
      }
    }
    if (type == -1) {
      mint inv = mint(len).pow(MOD - 2);
      for (int i = 0; i < len; i++) a[i] *= inv;
    }
  }
};
#line 1 "src/misc/read.hpp"
#define GC ch = getchar_unlocked()
ll read() {
  ll x = 0, f = 1, GC;
  while (ch < '0' || ch > '9') ch == '-' ? f = -1, GC : GC;
  while (ch >= '0' && ch <= '9') x = x * 10 + ch - '0', GC;
  return x * f;
}
#undef GC
#line 9 "test/math/ntt.test.cpp"

int main() {
  cin.tie(nullptr)->sync_with_stdio(false);
  int n = read(), m = read();
  Poly F, G;
  F.data.resize(n);
  for (int i = 0; i < n; i++) F.data[i] = read();
  G.data.resize(m);
  for (int i = 0; i < m; i++) G.data[i] = read();
  Poly H = F * G;
  for (int i = 0; i < n + m - 1; i++) {
    cout << H.data[i].data << " \n"[i == n + m - 2];
  }
  return 0;
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	all_same_00	AC	218 ms	20 MB
g++	all_same_01	AC	232 ms	20 MB
g++	all_same_02	AC	245 ms	20 MB
g++	all_same_03	AC	239 ms	20 MB
g++	example_00	AC	5 ms	4 MB
g++	example_01	AC	4 ms	4 MB
g++	fft_killer_00	AC	229 ms	20 MB
g++	fft_killer_01	AC	240 ms	20 MB
g++	fft_killer_02	AC	240 ms	20 MB
g++	fft_killer_03	AC	237 ms	20 MB
g++	fft_killer_04	AC	233 ms	20 MB
g++	fft_killer_05	AC	232 ms	20 MB
g++	fft_killer_06	AC	238 ms	20 MB
g++	fft_killer_07	AC	231 ms	20 MB
g++	fft_killer_08	AC	248 ms	20 MB
g++	fft_killer_09	AC	233 ms	20 MB
g++	max_ans_zero_00	AC	232 ms	20 MB
g++	max_random_00	AC	233 ms	20 MB
g++	max_random_01	AC	239 ms	20 MB
g++	medium_00	AC	8 ms	4 MB
g++	medium_01	AC	6 ms	4 MB
g++	medium_02	AC	7 ms	4 MB
g++	medium_all_zero_00	AC	7 ms	4 MB
g++	medium_pre_suf_zero_00	AC	5 ms	4 MB
g++	medium_pre_suf_zero_01	AC	5 ms	4 MB
g++	medium_pre_suf_zero_02	AC	5 ms	4 MB
g++	medium_pre_suf_zero_03	AC	5 ms	4 MB
g++	medium_pre_suf_zero_04	AC	5 ms	4 MB
g++	random_00	AC	220 ms	19 MB
g++	random_01	AC	225 ms	19 MB
g++	random_02	AC	113 ms	12 MB
g++	signed_overflow_00	AC	6 ms	4 MB
g++	small_00	AC	5 ms	4 MB
g++	small_01	AC	5 ms	4 MB
g++	small_02	AC	5 ms	4 MB
g++	small_03	AC	5 ms	4 MB
g++	small_04	AC	5 ms	4 MB
g++	small_05	AC	5 ms	4 MB
g++	small_06	AC	5 ms	4 MB
g++	small_07	AC	5 ms	4 MB
g++	small_08	AC	5 ms	4 MB
g++	small_09	AC	5 ms	4 MB
g++	small_10	AC	5 ms	4 MB
g++	small_11	AC	5 ms	4 MB
g++	small_12	AC	5 ms	4 MB
g++	small_13	AC	5 ms	4 MB
g++	small_14	AC	5 ms	4 MB
g++	small_15	AC	5 ms	4 MB
g++	small_and_large_00	AC	199 ms	18 MB
g++	small_and_large_01	AC	199 ms	18 MB
g++	small_and_large_02	AC	202 ms	18 MB
g++	small_and_large_03	AC	200 ms	18 MB
g++	unsigned_overflow_00	AC	5 ms	4 MB

Back to top page