Problem p02563 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>
using namespace std;

template <int mod> struct NumberTheoreticTransform {

  vector<int> rev, rts;
  int base, max_base, root;

  NumberTheoreticTransform() : base(1), rev{0, 1}, rts{0, 1} {
    assert(mod >= 3 && mod % 2 == 1);
    auto tmp = mod - 1;
    max_base = 0;
    while (tmp % 2 == 0)
      tmp >>= 1, max_base++;
    root = 2;
    while (mod_pow(root, (mod - 1) >> 1) == 1)
      ++root;
    assert(mod_pow(root, mod - 1) == 1);
    root = mod_pow(root, (mod - 1) >> max_base);
  }

  inline int mod_pow(int x, int n) {
    int ret = 1;
    while (n > 0) {
      if (n & 1)
        ret = mul(ret, x);
      x = mul(x, x);
      n >>= 1;
    }
    return ret;
  }

  inline int inverse(int x) { return mod_pow(x, mod - 2); }

  inline unsigned add(unsigned x, unsigned y) {
    x += y;
    if (x >= mod)
      x -= mod;
    return x;
  }

  inline unsigned mul(unsigned a, unsigned b) {
    return 1ull * a * b % (unsigned long long)mod;
  }

  void ensure_base(int nbase) {
    if (nbase <= base)
      return;
    rev.resize(1 << nbase);
    rts.resize(1 << nbase);
    for (int i = 0; i < (1 << nbase); i++) {
      rev[i] = (rev[i >> 1] >> 1) + ((i & 1) << (nbase - 1));
    }
    assert(nbase <= max_base);
    while (base < nbase) {
      int z = mod_pow(root, 1 << (max_base - 1 - base));
      for (int i = 1 << (base - 1); i < (1 << base); i++) {
        rts[i << 1] = rts[i];
        rts[(i << 1) + 1] = mul(rts[i], z);
      }
      ++base;
    }
  }

  void ntt(vector<int> &a) {
    const int n = (int)a.size();
    assert((n & (n - 1)) == 0);
    int zeros = __builtin_ctz(n);
    ensure_base(zeros);
    int shift = base - zeros;
    for (int i = 0; i < n; i++) {
      if (i < (rev[i] >> shift)) {
        swap(a[i], a[rev[i] >> shift]);
      }
    }
    for (int k = 1; k < n; k <<= 1) {
      for (int i = 0; i < n; i += 2 * k) {
        for (int j = 0; j < k; j++) {
          int z = mul(a[i + j + k], rts[j + k]);
          a[i + j + k] = add(a[i + j], mod - z);
          a[i + j] = add(a[i + j], z);
        }
      }
    }
  }

  vector<int> multiply(vector<int> a, vector<int> b) {
    int need = a.size() + b.size() - 1;
    int nbase = 1;
    while ((1 << nbase) < need)
      nbase++;
    ensure_base(nbase);
    int sz = 1 << nbase;
    a.resize(sz, 0);
    b.resize(sz, 0);
    ntt(a);
    ntt(b);
    int inv_sz = inverse(sz);
    for (int i = 0; i < sz; i++) {
      a[i] = mul(a[i], mul(b[i], inv_sz));
    }
    reverse(a.begin() + 1, a.end());
    ntt(a);
    a.resize(need);
    return a;
  }
};

int main() {
  int n, m;
  cin >> n >> m;
  vector<int> a(n), b(m);
  for (auto &x : a)
    cin >> x;
  for (auto &x : b)
    cin >> x;

  NumberTheoreticTransform<998244353> ntt;
  vector<int> c = ntt.multiply(a, b);
  for (auto x : c)
    cout << x << ' ';
  cout << '\\n';
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define for_range(i, a, b) for (int i = (a); i < (b); i++)
#define for_range_rev(i, a, b) for (int i = (b)-1; i >= (a); i++)
using namespace std;
using i64 = int_fast64_t;
using u64 = uint_fast64_t;
using u32 = uint_fast32_t;
// const u64 MOD = 1'000'000'007;
const u64 MOD = 998'244'353;
const int DX4[] = {1, 0, -1, 0};
const int DY4[] = {0, 1, 0, -1};
const int DX8[] = {0, 1, 1, 1, 0, -1, -1, -1};
const int DY8[] = {1, 1, 0, -1, -1, -1, 0, 1};

int ceil_log2(int n) {
  int x = 0;
  while ((1u << x) < (u32)(n))
    x++;
  return x;
}

struct mint {
  u64 a;
  static const u64 m = MOD;
  // static u64 m;

  // constexpr, noexcept などの意味：不明
  constexpr mint(const u64 x = 0) noexcept : a(x % m) {}
  constexpr u64 &val() noexcept { return a; }
  constexpr const u64 &val() const noexcept { return a; }
  constexpr mint operator+(const mint r) const noexcept {
    return mint(*this) += r;
  }
  constexpr mint operator-(const mint r) const noexcept {
    return mint(*this) -= r;
  }
  constexpr mint operator*(const mint r) const noexcept {
    return mint(*this) *= r;
  }
  constexpr mint operator/(const mint r) const noexcept {
    return mint(*this) /= r;
  }
  constexpr mint &operator+=(const mint r) noexcept {
    a += r.a;
    if (a >= m) {
      a -= m;
    }
    return *this;
  }
  constexpr mint &operator-=(const mint r) noexcept {
    if (a < r.a) {
      a += m;
    }
    a -= r.a;
    return *this;
  }
  constexpr mint &operator*=(const mint r) noexcept {
    a = a * r.a % m;
    return *this;
  }
  constexpr mint &operator/=(mint r) noexcept { return *this *= r.inv(); }
  constexpr mint pow(i64 e) noexcept {
    assert(0 <= e);
    mint x = *this, r = 1;
    while (e) {
      if (e & 1)
        r *= x;
      x *= x;
      e >>= 1;
    }
    return r;
  }
  constexpr mint inv() noexcept { return (*this).pow(m - 2); }
};

void butterfly(vector<mint> &a) {
  const int G = 3;
  int n = a.size();
  int h = ceil_log2(n);

  static bool f = true;
  static mint s[30];
  if (f) {
    f = false;
    mint g[30], b[30];
    int c = __builtin_ctz(MOD - 1);
    mint d = mint(G).pow((MOD - 1) >> c), e = d.inv();
    for_range_rev(i, 2, c + 1) {
      g[i - 2] = d;
      b[i - 2] = e;
      d *= d;
      e *= e;
    }
    mint m = 1;
    for_range(i, 0, c - 2) {
      s[i] = g[i] * m;
      m *= b[i];
    }
  }
  for_range(q, 1, h + 1) {
    int w = 1 << (q - 1), p = 1 << (h - q);
    mint m = 1;
    for_range(t, 0, w) {
      int o = t << (h - q + 1);
      for_range(i, 0, p) {
        auto l = a[i + o];
        auto r = a[i + o + p] * m;
        a[i + o] = l + r;
        a[i + o + p] = l - r;
      }
      m *= s[__builtin_ctz(~(u32)t)];
    }
  }
}

// ほぼ上のコピペ
void butterfly_inv(vector<mint> &a) {
  const int G = 3;
  int n = a.size();
  int h = ceil_log2(n);

  static bool f = true;
  static mint s[30];
  if (f) {
    f = false;
    mint g[30], b[30];
    int c = __builtin_ctz(MOD - 1);
    mint d = mint(G).pow((MOD - 1) >> c), e = d.inv();
    for_range_rev(i, 2, c + 1) {
      g[i - 2] = d;
      b[i - 2] = e;
      d *= d;
      e *= e;
    }
    mint m = 1;
    for_range(i, 0, c - 2) {
      s[i] = g[i] * m;
      m *= b[i];
    }
  }
  for_range(q, 1, h + 1) {
    int w = 1 << (q - 1), p = 1 << (h - q);
    mint m = 1;
    for_range(t, 0, w) {
      int o = t << (h - q + 1);
      for_range(i, 0, p) {
        auto l = a[i + o];
        // こっから
        auto r = a[i + o + p];
        a[i + o] = l + r;
        a[i + o + p] = (u64)(MOD + l.val() - r.val()) * m.val();
        // ここまで
      }
      m *= s[__builtin_ctz(~(u32)t)];
    }
  }
}

vector<mint> convolution(vector<mint> a, vector<mint> b) {
  int n = a.size(), m = b.size();
  if (!n || !m)
    return {};
  if (min(n, m) <= 60) {
    if (n < m) {
      swap(n, m);
      swap(a, b);
    }
    vector<mint> ans(n + m - 1);
    for_range(i, 0, n) for_range(j, 0, m) ans[i + j] += a[i] * b[j];
    return ans;
  }
  int z = 1 << ceil_log2(n + m - 1);
  a.resize(z);
  butterfly(a);
  b.resize(z);
  butterfly(b);
  for_range(i, 0, z) a[i] *= b[i];
  butterfly_inv(a);
  a.resize(n + m - 1);
  mint r = mint(z).inv();
  for_range(i, 0, n + m - 1) a[i] *= r;
  return a;
}

int main() {
  int n, m;
  cin >> n >> m;
  vector<mint> a(n), b(m);
  for (auto &x : a) {
    i64 t;
    cin >> t;
    x = mint(t);
  }
  for (auto &x : b) {
    i64 t;
    cin >> t;
    x = mint(t);
  }
  vector<mint> c = convolution(a, b);
  for_range(i, 0, n + m - 2) cout << c[i].val() << ' ';
  cout << c[n + m - 2].val() << '\\n';
}