代码

// v11: AVX2 Montgomery NTT, sparse u64 (4 lanes), contiguous per-stage twiddles.
#pragma GCC target("avx2")
#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <immintrin.h>

typedef unsigned long long u64;
typedef unsigned int u32;

const u32 MOD = 998244353u;
const u32 NINV = 998244351u;
const u32 R2 = 932051910u;
const u32 G = 3u;
const int MAXN = 1 << 18;

alignas(32) static u64 a[MAXN], b[MAXN];
alignas(32) static u64 roots[MAXN];
alignas(32) static u64 tw_fwd[MAXN];
alignas(32) static u64 tw_inv[MAXN];

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

static inline u32 mont_s(u32 a, u32 b) {
    u64 t = (u64)a * b;
    u32 m = (u32)t * NINV;
    u32 u = (u32)((t + (u64)m * MOD) >> 32);
    if (u >= MOD) u -= MOD;
    return u;
}

static inline __m256i mont4(__m256i x, __m256i y, const __m256i &ninv, const __m256i &modv, const __m256i &modm1) {
    __m256i t = _mm256_mul_epu32(x, y);
    __m256i m = _mm256_mul_epu32(t, ninv);
    __m256i mt = _mm256_mul_epu32(m, modv);
    __m256i r = _mm256_srli_epi64(_mm256_add_epi64(t, mt), 32);
    __m256i mask = _mm256_cmpgt_epi64(r, modm1);
    r = _mm256_sub_epi64(r, _mm256_and_si256(mask, modv));
    return r;
}

static void ntt_fwd(u64 *x, int n, const u64 *tw) {
    const __m256i ninv = _mm256_set1_epi64x(NINV);
    const __m256i modv = _mm256_set1_epi64x(MOD);
    const __m256i modm1 = _mm256_set1_epi64x(MOD - 1);
    for (int len = n; len > 1; len >>= 1) {
        int half = len >> 1;
        for (int i = 0; i < n; i += len) {
            u64 *y = x + i;
            int j = 0;
            for (; j + 4 <= half; j += 4) {
                __m256i u = _mm256_load_si256((__m256i*)(y + j));
                __m256i v = _mm256_load_si256((__m256i*)(y + j + half));
                __m256i s = _mm256_add_epi64(u, v);
                __m256i smask = _mm256_cmpgt_epi64(s, modm1);
                s = _mm256_sub_epi64(s, _mm256_and_si256(smask, modv));
                __m256i d = _mm256_sub_epi64(u, v);
                __m256i negmask = _mm256_cmpgt_epi64(_mm256_setzero_si256(), d);
                d = _mm256_add_epi64(d, _mm256_and_si256(negmask, modv));
                __m256i twv = _mm256_load_si256((__m256i*)(tw + j));
                __m256i dv = mont4(d, twv, ninv, modv, modm1);
                _mm256_store_si256((__m256i*)(y + j), s);
                _mm256_store_si256((__m256i*)(y + j + half), dv);
            }
            for (; j < half; j++) {
                u64 u = y[j];
                u64 v = y[j + half];
                u64 s = u + v; if (s >= MOD) s -= MOD;
                u64 d = u - v; if (d >= MOD) d += MOD;
                y[j] = s;
                y[j + half] = mont_s((u32)d, (u32)tw[j]);
            }
        }
        tw += half;
    }
}

static void ntt_inv(u64 *x, int n, const u64 *tw) {
    const __m256i ninv = _mm256_set1_epi64x(NINV);
    const __m256i modv = _mm256_set1_epi64x(MOD);
    const __m256i modm1 = _mm256_set1_epi64x(MOD - 1);
    for (int len = 2; len <= n; len <<= 1) {
        int half = len >> 1;
        for (int i = 0; i < n; i += len) {
            u64 *y = x + i;
            int j = 0;
            for (; j + 4 <= half; j += 4) {
                __m256i u = _mm256_load_si256((__m256i*)(y + j));
                __m256i v = _mm256_load_si256((__m256i*)(y + j + half));
                __m256i twv = _mm256_load_si256((__m256i*)(tw + j));
                v = mont4(v, twv, ninv, modv, modm1);
                __m256i s = _mm256_add_epi64(u, v);
                __m256i smask = _mm256_cmpgt_epi64(s, modm1);
                s = _mm256_sub_epi64(s, _mm256_and_si256(smask, modv));
                __m256i d = _mm256_sub_epi64(u, v);
                __m256i negmask = _mm256_cmpgt_epi64(_mm256_setzero_si256(), d);
                d = _mm256_add_epi64(d, _mm256_and_si256(negmask, modv));
                _mm256_store_si256((__m256i*)(y + j), s);
                _mm256_store_si256((__m256i*)(y + j + half), d);
            }
            for (; j < half; j++) {
                u64 u = y[j];
                u64 v = mont_s((u32)y[j + half], (u32)tw[j]);
                u64 s = u + v; if (s >= MOD) s -= MOD;
                u64 d = u - v; if (d >= MOD) d += MOD;
                y[j] = s;
                y[j + half] = d;
            }
        }
        tw += half;
    }
}

// fast input buffer
static const int BUFSZ = 1 << 20;
static char inbuf[BUFSZ];
static size_t inpos = 0, inlen = 0;
static inline int readbyte() {
    if (inpos >= inlen) {
        inlen = fread(inbuf, 1, BUFSZ, stdin);
        inpos = 0;
        if (inlen == 0) return -1;
    }
    return (unsigned char)inbuf[inpos++];
}
static inline int readint() {
    int c = readbyte();
    while (c == ' ' || c == '\n' || c == '\r' || c == '\t') c = readbyte();
    int x = 0;
    while (c >= '0' && c <= '9') { x = x * 10 + (c - '0'); c = readbyte(); }
    return x;
}

static char outbuf[1 << 21];
static size_t outpos = 0;
static inline void putc(char c) { outbuf[outpos++] = c; }
static inline void putint(int x) {
    if (x == 0) { putc('0'); return; }
    char tmp[12]; int t = 0;
    while (x) { tmp[t++] = '0' + (x % 10); x /= 10; }
    while (t) putc(tmp[--t]);
}

int main() {
    int n = readint();
    int m = readint();
    int na = n + 1;
    int nb = m + 1;
    for (int i = 0; i < na; i++) a[i] = (u32)readint();
    for (int i = 0; i < nb; i++) b[i] = (u32)readint();

    int size = 1;
    while (size < na + nb - 1) size <<= 1;

    u32 wstep = (MOD - 1) / (u32)size;
    u32 w = modpow(G, wstep);
    u32 invw = modpow(w, MOD - 2);
    u64 cur = 1;
    for (int i = 0; i < size; i++) { roots[i] = mont_s((u32)cur, R2); cur = cur * w % MOD; }
    {
        u64 *p = tw_fwd;
        for (int len = size; len > 1; len >>= 1) {
            int half = len >> 1;
            int step = size / len;
            for (int j = 0; j < half; j++) p[j] = roots[j * step];
            p += half;
        }
    }
    cur = 1;
    for (int i = 0; i < size; i++) { roots[i] = mont_s((u32)cur, R2); cur = cur * invw % MOD; }
    {
        u64 *p = tw_inv;
        for (int len = 2; len <= size; len <<= 1) {
            int half = len >> 1;
            int step = size / len;
            for (int j = 0; j < half; j++) p[j] = roots[j * step];
            p += half;
        }
    }

    for (int i = 0; i < size; i++) a[i] = mont_s((u32)a[i], R2);
    for (int i = 0; i < size; i++) b[i] = mont_s((u32)b[i], R2);

    ntt_fwd(a, size, tw_fwd);
    ntt_fwd(b, size, tw_fwd);

    {
        const __m256i ninv = _mm256_set1_epi64x(NINV);
        const __m256i modv = _mm256_set1_epi64x(MOD);
        const __m256i modm1 = _mm256_set1_epi64x(MOD - 1);
        for (int i = 0; i < size; i += 4) {
            __m256i x = _mm256_load_si256((__m256i*)(a + i));
            __m256i y = _mm256_load_si256((__m256i*)(b + i));
            _mm256_store_si256((__m256i*)(a + i), mont4(x, y, ninv, modv, modm1));
        }
    }

    ntt_inv(a, size, tw_inv);

    for (int i = 0; i < size; i++) a[i] = mont_s((u32)a[i], 1);
    u32 ninv_scale = modpow((u32)size, MOD - 2);
    for (int i = 0; i < size; i++) a[i] = (u64)((u32)((u64)a[i] * ninv_scale % MOD));

    int outn = n + m + 1;
    for (int i = 0; i < outn; i++) {
        if (i) putc(' ');
        putint((int)a[i]);
    }
    putc('\n');
    fwrite(outbuf, 1, outpos, stdout);
    return 0;
}

评测结果