NTT & FFT

From EOJ Wiki
Revision as of 11:56, 22 March 2018 by Ultmaster (talk | contribs) (Created page with "== NTT == <syntaxhighlight lang='cpp'> typedef long long LL; const int MAXN = 3e5 + 10; const int MOD = 998244353; const int G = 3; namespace NTT { int N, a[MAXN], b[MAX...")
(diff) ← Older revision | Latest revision (diff) | Newer revision → (diff)
Jump to navigation Jump to search

NTT

typedef long long LL;
const int MAXN = 3e5 + 10;
const int MOD = 998244353;
const int G = 3;

namespace NTT {
    int N, a[MAXN], b[MAXN];

    int pown(LL x, LL n) {
        LL ret = MOD != 1; x %= MOD;
        while (n) {
            if (n & 1) ret = ret * x % MOD;
            x = x * x % MOD;
            n >>= 1;
        }
        return int(ret);
    }

    void clear() {
        memset(a, 0, sizeof a);
        memset(b, 0, sizeof b);
    }

    void ntt(int * a, int N, int f) {
        int i, j = 0, t, k;
        for (i = 1; i < N - 1; i++) {
            for (t = N; j ^= t >>= 1, ~j & t;);
            if (i < j) {
                swap(a[i], a[j]);
            }
        }
        for (i = 1; i < N; i <<= 1) {
            t = i << 1;
            int wn = pown(G, (MOD - 1) / t);
            for (j = 0; j < N; j += t) {
                int w = 1;
                for (k = 0; k < i; k++, w = 1LL * w * wn % MOD) {
                    int x = a[j + k], y = 1LL * w * a[j + k + i] % MOD;
                    a[j + k] = (x + y) % MOD, a[j + k + i] = (x - y + MOD) % MOD;
                }
            }
        }
        if (f == -1) {
            reverse(a + 1, a + N);
            int inv = pown(N, MOD - 2);
            for (i = 0; i < N; i++)
                a[i] = 1LL * a[i] * inv % MOD;
        }
    }

    void conv() {
        ntt(a, N, 1);
        ntt(b, N, 1);
        for (int i = 0; i < N; ++i)
            a[i] = 1LL * a[i] * b[i] % MOD;
        ntt(a, N, -1);
    }
};