代码

// NTT mod 998244353, scalar, DIT + bitrev (correctness baseline)
typedef unsigned long long u64;
typedef unsigned u32;

const u32 MOD = 998244353;
const u32 G = 3;
const int MAXL = 1 << 21;

static u32 A[MAXL];
static u32 B[MAXL];
static u32 roots[MAXL];      // forward roots
static u32 roots_inv[MAXL];  // inverse roots

static inline u32 mul_mod(u32 x, u32 y) {
    return (u32)((u64)x * y % MOD);
}
static inline u32 add_mod(u32 x, u32 y) {
    u32 s = x + y;
    return s >= MOD ? s - MOD : s;
}
static inline u32 sub_mod(u32 x, u32 y) {
    return x >= y ? x - y : x + MOD - y;
}

static inline u32 modpow(u32 base, u64 e) {
    u64 r = 1, b = base;
    while (e) {
        if (e & 1) r = r * b % MOD;
        b = b * b % MOD;
        e >>= 1;
    }
    return (u32)r;
}

static void ntt(u32 *x, int n, const u32 *rts) {
    // bit reversal
    for (int i = 1, j = 0; i < n; i++) {
        int bit = n >> 1;
        for (; j & bit; bit >>= 1) j ^= bit;
        j ^= bit;
        if (i < j) { u32 t = x[i]; x[i] = x[j]; x[j] = t; }
    }
    for (int len = 2; len <= n; len <<= 1) {
        int half = len >> 1;
        int step = n / len;
        for (int i = 0; i < n; i += len) {
            u32 *y = x + i;
            for (int j = 0; j < half; j++) {
                u32 u = y[j];
                u32 v = mul_mod(y[j + half], rts[j * step]);
                y[j] = add_mod(u, v);
                y[j + half] = sub_mod(u, v);
            }
        }
    }
}

void poly_multiply(unsigned *a, int n, unsigned *b, int m, unsigned *c) {
    int L = 1;
    while (L < n + m + 2) L <<= 1;
    // copy + zero pad
    for (int i = 0; i <= n; i++) A[i] = a[i];
    for (int i = 0; i <= m; i++) B[i] = b[i];
    for (int i = n + 1; i < L; i++) A[i] = 0;
    for (int i = m + 1; i < L; i++) B[i] = 0;

    // precompute roots
    u32 w = modpow(G, (MOD - 1) / L);
    u32 wi = modpow(w, MOD - 2);
    u64 cur = 1;
    for (int i = 0; i < L; i++) { roots[i] = (u32)cur; cur = cur * w % MOD; }
    cur = 1;
    for (int i = 0; i < L; i++) { roots_inv[i] = (u32)cur; cur = cur * wi % MOD; }

    ntt(A, L, roots);
    ntt(B, L, roots);
    for (int i = 0; i < L; i++) A[i] = mul_mod(A[i], B[i]);
    ntt(A, L, roots_inv);

    u32 ninv = modpow(L, MOD - 2);
    for (int i = 0; i <= n + m; i++) c[i] = mul_mod(A[i], ninv);
}

评测结果