test/FPS/FFT_fast.test.cpp

View this file on GitHub
Last update: 2024-05-31 19:00:40+09:00
Problem: https://judge.yosupo.jp/problem/convolution_mod

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/convolution_mod"
#include "test/template.hpp"
#include "src/extra/modint_fast.hpp"
#include "src/FPS/FFT_fast.hpp"

int main() {
   cin.tie(0)->sync_with_stdio(0);
   ll N, M;
   cin >> N >> M;
   vector<mm> A(N), B(M);
   for(mm& a : A) cin >> a.x;
   for(mm& b : B) cin >> b.x;

   auto C = conv(move(A), move(B));
   rep(i, 0, sz(C)) cout << C[i].x << " \n"[i + 1 == sz(C)];
}

#line 1 "test/FPS/FFT_fast.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/convolution_mod"
#line 1 "test/template.hpp"
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const ll INF = LLONG_MAX / 4;
#define rep(i, a, b) for(ll i = a; i < (b); i++)
#define all(a) begin(a), end(a)
#define sz(a) ssize(a)
bool chmin(auto& a, auto b) { return a > b ? a = b, 1 : 0; }
bool chmax(auto& a, auto b) { return a < b ? a = b, 1 : 0; }
#line 1 "src/extra/modint_fast.hpp"
const uint32_t mod = 998244353;
struct mm {
   uint32_t x;
   mm() : x(0) {}
   template<class T> mm(T x_) : x(x_ % mod) {
      if(x >= mod) x += mod;
   }
   friend mm operator+(mm a, mm b) {
      a.x += b.x;
      if(a.x >= mod) a.x -= mod;
      return a;
   }
   friend mm operator-(mm a, mm b) {
      a.x -= b.x;
      if(a.x >= mod) a.x += mod;
      return a;
   }
   friend mm operator*(mm a, mm b) { return (uint64_t)a.x * b.x; }
   friend mm operator/(mm a, mm b) { return a * b.inv(); }
   friend mm& operator+=(mm& a, mm b) { return a = a + b; }
   friend mm& operator-=(mm& a, mm b) { return a = a - b; }
   friend mm& operator*=(mm& a, mm b) { return a = a * b; }
   friend mm& operator/=(mm& a, mm b) { return a = a * b.inv(); }
   mm inv() const { return pow(mod - 2); }
   mm pow(ll b) const {
      mm a = *this, c = 1;
      while(b) {
         if(b & 1) c *= a;
         a *= a;
         b >>= 1;
      }
      return c;
   }
};
#line 1 "src/FPS/FFT_fast.hpp"
// modint を u32 にして加減算を真面目にやると速い
mm g = 3;  // 原始根
void fft(vector<mm>& a) {
   ll n = sz(a), lg = __lg(n);
   static auto z = [] {
      vector<mm> z(30);
      mm s = 1;
      rep(i, 2, 32) {
         z[i - 2] = s * g.pow(mod >> i);
         s *= g.inv().pow(mod >> i);
      }
      return z;
   }();
   rep(l, 0, lg) {
      ll w = 1 << (lg - l - 1);
      mm s = 1;
      rep(k, 0, 1 << l) {
         ll o = k << (lg - l);
         rep(i, o, o + w) {
            mm x = a[i], y = a[i + w] * s;
            a[i] = x + y;
            a[i + w] = x - y;
         }
         s *= z[countr_zero<uint64_t>(~k)];
      }
   }
}
// コピペ
void ifft(vector<mm>& a) {
   ll n = sz(a), lg = __lg(n);
   static auto z = [] {
      vector<mm> z(30);
      mm s = 1;
      rep(i, 2, 32) {  // g を逆数に
         z[i - 2] = s * g.inv().pow(mod >> i);
         s *= g.pow(mod >> i);
      }
      return z;
   }();
   for(ll l = lg; l--;) {  // 逆順に
      ll w = 1 << (lg - l - 1);
      mm s = 1;
      rep(k, 0, 1 << l) {
         ll o = k << (lg - l);
         rep(i, o, o + w) {
            mm x = a[i], y = a[i + w];  // *s を下に移動
            a[i] = x + y;
            a[i + w] = (x - y) * s;
         }
         s *= z[countr_zero<uint64_t>(~k)];
      }
   }
}
vector<mm> conv(vector<mm> a, vector<mm> b) {
   if(a.empty() || b.empty()) return {};
   size_t s = sz(a) + sz(b) - 1, n = bit_ceil(s);
   // if(min(sz(a), sz(b)) <= 60) 愚直に掛け算
   a.resize(n);
   b.resize(n);
   fft(a);
   fft(b);
   mm inv = mm(n).inv();
   rep(i, 0, n) a[i] *= b[i] * inv;
   ifft(a);
   a.resize(s);
   return a;
}
#line 5 "test/FPS/FFT_fast.test.cpp"

int main() {
   cin.tie(0)->sync_with_stdio(0);
   ll N, M;
   cin >> N >> M;
   vector<mm> A(N), B(M);
   for(mm& a : A) cin >> a.x;
   for(mm& b : B) cin >> b.x;

   auto C = conv(move(A), move(B));
   rep(i, 0, sz(C)) cout << C[i].x << " \n"[i + 1 == sz(C)];
}