bingxuan9112

NTT

Mar 5th, 2020
245
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 1.27 KB | None | 0 0
  1. #include <bits/stdc++.h>
  2.  
  3. using namespace std;
  4.  
  5. int64_t modpow(int64_t e,int64_t p,int64_t m) {
  6.     int64_t r = 1;
  7.     while(p) (p&1)&&(r=r*e%m), e=e*e%m, p>>=1;
  8.     return r;
  9. }
  10. void NTT(int64_t F[],int n,bool inv,int64_t m=998244353,int64_t g=3) {
  11.     for(int i = 0, j = 0; i < n; i++) {
  12.         if(i < j) swap(F[i],F[j]);
  13.         for(int k = n>>1; (j^=k) < k; k>>=1);
  14.     }
  15.     for(int step = 1; step < n; step *= 2) {
  16.         int64_t root = modpow(g,(m-1)/(step*2),m);
  17.         if(inv) root = modpow(root,m-2,m);
  18.         for(int i = 0; i < n; i += step*2) {
  19.             int64_t now = 1;
  20.             for(int j = 0; j < step; j++) {
  21.                 int64_t a = F[i+j];
  22.                 int64_t b = F[i+j+step]*now%m;
  23.                 F[i+j] = (a+b<m ? a+b : a+b-m);
  24.                 F[i+j+step] = (a-b<0 ? m+a-b : a-b);
  25.                 now = now*root%m;
  26.             }
  27.         }
  28.     }
  29.     if(inv) {
  30.         int64_t in = modpow(n,m-2,m);
  31.         for(int i = 0; i < n; i++) F[i]=F[i]*in%m;
  32.     }
  33. }
  34.  
  35. int64_t A[1<<20] = {1,2}, B[1<<20] = {1,2,3}; // (1+2x) * (1+2x+3x^2)
  36. signed main() {
  37.     // A*B
  38.     int n = 1<<__lg(2+3)+1; // n = 2^k >= deg(A)+deg(B)
  39.     NTT(A,n,0);
  40.     NTT(B,n,0);
  41.     for(int i = 0; i < n; i++) A[i] = A[i]*B[i]%998244353;
  42.     NTT(A,n,1);
  43. }
Advertisement
Add Comment
Please, Sign In to add comment