代码

// NTT mod 998244353, radix-4 DIF forward + DIT inverse (no bitrev), scalar Montgomery
typedef unsigned long long u64;
typedef unsigned u32;

const u32 MOD = 998244353u;
const u32 NINV = 998244351u;
const u32 ONE  = 301989884u;
const u32 R2   = 932051910u;
const u32 IROOT = 911660635u;
const int MAXL = 1 << 21;

static u32 A[MAXL];
static u32 B[MAXL];
static u32 roots[MAXL];
static u32 roots_inv[MAXL];

static inline u32 mont_mul(u32 x, u32 y) {
    u64 t = (u64)x * y;
    u32 m = (u32)t * NINV;
    u64 u = (t + (u64)m * MOD) >> 32;
    if (u >= MOD) u -= MOD;
    return (u32)u;
}
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 to_mont(u32 x) { return mont_mul(x, R2); }
static u32 mont_pow(u32 base, u64 e) {
    u32 r = ONE, b = base;
    while (e) { if (e & 1) r = mont_mul(r, b); b = mont_mul(b, b); e >>= 1; }
    return r;
}

// n must be a power of two. If n = 4^k, pure radix-4. If n = 2*4^k, radix-2 top + radix-4.
static void ntt_fwd(u32 *x, int n, const u32 *rts, u32 iroot) {
    int start_len;
    if (n % 3 == 2) {  // n = 2 * 4^k
        int half = n >> 1;
        for (int j = 0; j < half; j++) {
            u32 u = x[j], v = x[j + half];
            x[j] = add_mod(u, v);
            x[j + half] = mont_mul(sub_mod(u, v), rts[j]);
        }
        start_len = n >> 1;
    } else {
        start_len = n;  // n = 4^k
    }
    for (int len = start_len; len >= 4; len >>= 2) {
        int m = len >> 2;
        int step = n / len;
        for (int i = 0; i < n; i += len) {
            u32 *y = x + i;
            for (int j = 0; j < m; j++) {
                u32 x0 = y[j], x1 = y[j + m], x2 = y[j + 2*m], x3 = y[j + 3*m];
                u32 a0 = add_mod(x0, x2);
                u32 a1 = sub_mod(x0, x2);
                u32 a2 = add_mod(x1, x3);
                u32 a3 = sub_mod(x1, x3);
                u32 ia3 = mont_mul(a3, iroot);
                u32 w1 = rts[j * step];
                u32 w2 = rts[2 * j * step];
                u32 w3 = rts[3 * j * step];
                y[j]       = add_mod(a0, a2);
                y[j + m]   = mont_mul(add_mod(a1, ia3), w1);
                y[j + 2*m] = mont_mul(sub_mod(a0, a2), w2);
                y[j + 3*m] = mont_mul(sub_mod(a1, ia3), w3);
            }
        }
    }
}

static void ntt_inv(u32 *x, int n, const u32 *rts, u32 iroot) {
    int end_len;
    if (n % 3 == 2) {
        end_len = n >> 1;
    } else {
        end_len = n;
    }
    for (int len = 4; len <= end_len; len <<= 2) {
        int m = len >> 2;
        int step = n / len;
        for (int i = 0; i < n; i += len) {
            u32 *y = x + i;
            for (int j = 0; j < m; j++) {
                u32 w1 = rts[j * step];
                u32 w2 = rts[2 * j * step];
                u32 w3 = rts[3 * j * step];
                u32 y0 = y[j];
                u32 y1 = mont_mul(y[j + m], w1);
                u32 y2 = mont_mul(y[j + 2*m], w2);
                u32 y3 = mont_mul(y[j + 3*m], w3);
                u32 b0 = add_mod(y0, y2);
                u32 b1 = sub_mod(y0, y2);
                u32 b2 = add_mod(y1, y3);
                u32 b3 = sub_mod(y1, y3);
                u32 ib3 = mont_mul(b3, iroot);
                y[j]       = add_mod(b0, b2);
                y[j + m]   = add_mod(b1, ib3);
                y[j + 2*m] = sub_mod(b0, b2);
                y[j + 3*m] = sub_mod(b1, ib3);
            }
        }
    }
    if (n % 3 == 2) {
        int half = n >> 1;
        for (int j = 0; j < half; j++) {
            u32 u = x[j], v = mont_mul(x[j + half], rts[j]);
            x[j] = add_mod(u, v);
            x[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;
    for (int i = 0; i <= n; i++) A[i] = to_mont(a[i]);
    for (int i = 0; i <= m; i++) B[i] = to_mont(b[i]);
    for (int i = n + 1; i < L; i++) A[i] = 0;
    for (int i = m + 1; i < L; i++) B[i] = 0;

    u32 w = mont_pow(to_mont(3), (MOD - 1) / L);
    u32 wi = mont_pow(w, MOD - 2);
    u32 cur = ONE;
    for (int i = 0; i < L; i++) { roots[i] = cur; cur = mont_mul(cur, w); }
    cur = ONE;
    for (int i = 0; i < L; i++) { roots_inv[i] = cur; cur = mont_mul(cur, wi); }

    u32 iplus = to_mont(IROOT);
    u32 iminus = mont_mul(iplus, to_mont(MOD - 1));

    ntt_fwd(A, L, roots, iplus);
    ntt_fwd(B, L, roots, iplus);
    for (int i = 0; i < L; i++) A[i] = mont_mul(A[i], B[i]);
    ntt_inv(A, L, roots_inv, iminus);

    u32 linv_mont = mont_pow(to_mont((u32)(L % MOD)), MOD - 2);
    u32 linv_std = mont_mul(linv_mont, 1);
    for (int i = 0; i <= n + m; i++) c[i] = mont_mul(A[i], linv_std);
}

评测结果