提交记录 61672


用户 题目 状态 得分 用时 内存 语言 代码长度
saffah_dsh_v41_0919 1002. 测测你的多项式乘法 Compile Error 0 0 ns 0 KB C++17 21.65 KB
提交时间 评测时间
2026-09-19 22:23:50 2026-09-19 22:23:52
// 1002i: polynomial multiply.  AVX2 NTT mod 998244353 (DIF forward / DIT inverse,
// Shoup twiddles, batched 8-point tail kernel; inverse scaled inside its last stage).
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <immintrin.h>
#include <sys/auxv.h>
typedef unsigned long long U64;
struct DI { U64 abi; const char *in; U64 insz; char *out; U64 outlim; U64 outsz; char *err;
            U64 errlim; U64 errsz; const char *IB; U64 IBlim; char *OB; U64 OBlim; U64 tscfreq; }
            __attribute__((packed));

// Fast AVX2 NTT mod 998244353.
//   fwd(a,n): natural order -> bit-reversed order  (DIF, vectorized over the twiddle index)
//   inv(a,n): bit-reversed  -> natural order       (DIT, same but conjugated)
// The last three DIF stages (len = 8,4,2) are fused into one branch-free 8-point
// kernel per contiguous group of 8, which removes the scalar tail that otherwise
// dominates small sizes.  Pointwise multiplication of two spectra is valid because
// both are in the same bit-reversed order and the inverse consumes that order.

typedef unsigned u32;
typedef unsigned long long u64;
static const u32 MOD = 998244353;
#ifndef POLY_N1BITS
#define POLY_N1BITS 11
#endif
#ifndef POLY_N2BITS
#define POLY_N2BITS 10
#endif

#define AVX2FN __attribute__((target("avx2")))

static inline u32 fpow(u32 a, u32 e) {
    u64 r = 1, b = a;
    while (e) { if (e & 1) r = r * b % MOD; b = b * b % MOD; e >>= 1; }
    return (u32)r;
}
AVX2FN static inline __m256i mulhi8(__m256i a, __m256i b) {
    __m256i lo = _mm256_mul_epu32(a, b);
    __m256i a1 = _mm256_srli_epi64(a, 32), b1 = _mm256_srli_epi64(b, 32);
    __m256i hi = _mm256_mul_epu32(a1, b1);
    return _mm256_blend_epi32(_mm256_srli_epi64(lo, 32), hi, 0xAA);
}
AVX2FN static inline __m256i shoup8(__m256i a, __m256i w, __m256i ws) {
    __m256i q = mulhi8(a, ws);
    __m256i r = _mm256_sub_epi32(_mm256_mullo_epi32(a, w), _mm256_mullo_epi32(q, _mm256_set1_epi32((int)MOD)));
    __m256i ge = _mm256_cmpgt_epi32(r, _mm256_set1_epi32((int)MOD - 1));
    return _mm256_sub_epi32(r, _mm256_and_si256(ge, _mm256_set1_epi32((int)MOD)));
}
AVX2FN static inline __m256i addmod8(__m256i a, __m256i b) {
    __m256i s = _mm256_add_epi32(a, b);
    return _mm256_sub_epi32(s, _mm256_and_si256(_mm256_cmpgt_epi32(s, _mm256_set1_epi32((int)MOD - 1)), _mm256_set1_epi32((int)MOD)));
}
AVX2FN static inline __m256i submod8(__m256i a, __m256i b) {
    __m256i lt = _mm256_cmpgt_epi32(b, a);
    return _mm256_add_epi32(_mm256_sub_epi32(a, b), _mm256_and_si256(lt, _mm256_set1_epi32((int)MOD)));
}
static inline u32 addm(u32 a, u32 b) { u32 s = a + b; return s >= MOD ? s - MOD : s; }
static inline u32 subm(u32 a, u32 b) { return a >= b ? a - b : a + MOD - b; }
static inline u32 mulm(u32 a, u32 b) { return (u32)((u64)a * b % MOD); }

struct Tables2 {
    int n = 0;
    u32 fw[1 << 19], fws[1 << 19], iw[1 << 19], iws[1 << 19];   // per-instance (T1 and T2 need different tables)
    int foff[40] = {0}, ioff[40] = {0};
    // build the twiddle sequence for stage len: values w_len^j, j = 0..half-1
    static void build(u32 *dst, int half, u32 wlen) {
        u32 wls = (u32)(((__uint128_t)wlen << 32) / MOD);        // Shoup constant for wlen
        u64 cur = 1;
        for (int j = 0; j < half; j++) {
            dst[j] = (u32)cur;
            u32 q = (u32)(((u64)cur * wls) >> 32);
            u32 r = (u32)((u64)cur * wlen - (u64)q * MOD);
            cur = r >= MOD ? r - MOD : r;
        }
    }
    // Build the forward and inverse twiddle chains together.  Each chain is a serial
    // modular-multiply dependency chain; interleaving two of them lets the out-of-order
    // core overlap the latencies, which is where most of the 16M-tick build cost goes.
    // Fused build: both twiddle chains run in lockstep (independent serial
    // modular-multiply chains -> ILP) and each entry's Shoup constant
    // ws = floor(w * 2^32 / MOD) is produced in the same iteration, so the extra
    // O(N) pass over both tables (and its memory traffic) disappears entirely.
    static void build2(u32 *d0, u32 *s0out, u32 *d1, u32 *s1out, int half, u32 w0, u32 w1) {
        u32 s0 = (u32)(((__uint128_t)w0 << 32) / MOD);
        u32 s1 = (u32)(((__uint128_t)w1 << 32) / MOD);
        const double inv = 4294967296.0 / (double)MOD;
        u64 c0 = 1, c1 = 1;
        for (int j = 0; j < half; j++) {
            {   // ws for the forward entry
                u32 x = (u32)c0;
                u32 q = (u32)((double)x * inv);
                u64 lhs = (u64)q * MOD, rhs = (u64)x << 32;
                if (lhs > rhs) q--; else if (lhs + MOD <= rhs) q++;
                s0out[j] = q;
            }
            d0[j] = (u32)c0;
            {   // ws for the inverse entry (independent -> overlaps the chain above)
                u32 x = (u32)c1;
                u32 q = (u32)((double)x * inv);
                u64 lhs = (u64)q * MOD, rhs = (u64)x << 32;
                if (lhs > rhs) q--; else if (lhs + MOD <= rhs) q++;
                s1out[j] = q;
            }
            d1[j] = (u32)c1;
            u32 qa = (u32)((c0 * s0) >> 32);
            u32 ra = (u32)(c0 * w0 - (u64)qa * MOD);
            c0 = ra >= MOD ? ra - MOD : ra;
            u32 qb = (u32)((c1 * s1) >> 32);
            u32 rb = (u32)(c1 * w1 - (u64)qb * MOD);
            c1 = rb >= MOD ? rb - MOD : rb;
        }
    }
    // ws = floor(w * 2^32 / MOD) using a double reciprocal plus exact correction
    static void buildws(u32 *dst, const u32 *w, int half) {
        const double inv = 4294967296.0 / (double)MOD;
        for (int j = 0; j < half; j++) {
            u32 x = w[j];
            u32 q = (u32)((double)x * inv);
            u64 lhs = (u64)q * MOD, rhs = (u64)x << 32;
            if (lhs > rhs) q--;
            else if (lhs + MOD <= rhs) q++;
            dst[j] = q;
        }
    }
    void init(int n_) {
        n = n_;
        if (n > (1 << 19)) n = 1 << 19;
        int idx = 0, idx2 = 0;
        for (int len = n; len >= 16; len >>= 1) {
            int b = __builtin_ctz(len);
            foff[b] = idx; ioff[b] = idx2;
            int half = len >> 1;
            build2(fw + idx, fws + idx, iw + idx2, iws + idx2, half,
                   fpow(3, (MOD - 1) / len),
                   fpow(3, MOD - 1 - (MOD - 1) / len));
            idx += half; idx2 += half;
        }
    }
};


AVX2FN static inline void tr8i(__m256i *r) {
    __m256i t0=_mm256_unpacklo_epi32(r[0],r[1]), t1=_mm256_unpackhi_epi32(r[0],r[1]);
    __m256i t2=_mm256_unpacklo_epi32(r[2],r[3]), t3=_mm256_unpackhi_epi32(r[2],r[3]);
    __m256i t4=_mm256_unpacklo_epi32(r[4],r[5]), t5=_mm256_unpackhi_epi32(r[4],r[5]);
    __m256i t6=_mm256_unpacklo_epi32(r[6],r[7]), t7=_mm256_unpackhi_epi32(r[6],r[7]);
    __m256i u0=_mm256_unpacklo_epi64(t0,t2), u1=_mm256_unpackhi_epi64(t0,t2);
    __m256i u2=_mm256_unpacklo_epi64(t1,t3), u3=_mm256_unpackhi_epi64(t1,t3);
    __m256i u4=_mm256_unpacklo_epi64(t4,t6), u5=_mm256_unpackhi_epi64(t4,t6);
    __m256i u6=_mm256_unpacklo_epi64(t5,t7), u7=_mm256_unpackhi_epi64(t5,t7);
    r[0]=_mm256_permute2x128_si256(u0,u4,0x20); r[1]=_mm256_permute2x128_si256(u1,u5,0x20);
    r[2]=_mm256_permute2x128_si256(u2,u6,0x20); r[3]=_mm256_permute2x128_si256(u3,u7,0x20);
    r[4]=_mm256_permute2x128_si256(u0,u4,0x31); r[5]=_mm256_permute2x128_si256(u1,u5,0x31);
    r[6]=_mm256_permute2x128_si256(u2,u6,0x31); r[7]=_mm256_permute2x128_si256(u3,u7,0x31);
}
static inline void k8fwd_scalar(u32 *x, const u32 *w8, const u32 *w4) {
    for (int j = 0; j < 4; j++) { u32 u = x[j], v = x[j + 4]; x[j] = addm(u, v); x[j + 4] = mulm(subm(u, v), w8[j]); }
    for (int b = 0; b < 8; b += 4) for (int j = 0; j < 2; j++) { u32 u = x[b + j], v = x[b + j + 2]; x[b + j] = addm(u, v); x[b + j + 2] = mulm(subm(u, v), w4[j]); }
    for (int b = 0; b < 8; b += 2) { u32 u = x[b], v = x[b + 1]; x[b] = addm(u, v); x[b + 1] = subm(u, v); }
}
static inline void k8inv_scalar(u32 *x, const u32 *w8i, const u32 *w4i) {
    for (int b = 0; b < 8; b += 2) { u32 u = x[b], v = x[b + 1]; x[b] = addm(u, v); x[b + 1] = subm(u, v); }
    for (int b = 0; b < 8; b += 4) for (int j = 0; j < 2; j++) { u32 u = x[b + j], v = mulm(x[b + j + 2], w4i[j]); x[b + j] = addm(u, v); x[b + j + 2] = subm(u, v); }
    for (int j = 0; j < 4; j++) { u32 u = x[j], v = mulm(x[j + 4], w8i[j]); x[j] = addm(u, v); x[j + 4] = subm(u, v); }
}
// generic scalar DIF (natural -> bit-reversed) for tiny n
static void small_fwd(u32 *a, int n) {
    for (int len = n; len >= 2; len >>= 1) {
        int half = len >> 1;
        u32 lw = fpow(3, (MOD - 1) / len);
        for (int i = 0; i < n; i += len) {
            u64 cur = 1;
            for (int j = 0; j < half; j++) {
                u32 u = a[i + j], v = a[i + j + half];
                a[i + j] = addm(u, v); a[i + j + half] = mulm(subm(u, v), (u32)cur);
                cur = cur * lw % MOD;
            }
        }
    }
}
static void small_inv(u32 *a, int n) {
    for (int len = 2; len <= n; len <<= 1) {
        int half = len >> 1;
        u32 lw = fpow(3, MOD - 1 - (MOD - 1) / len);
        for (int i = 0; i < n; i += len) {
            u64 cur = 1;
            for (int j = 0; j < half; j++) {
                u32 u = a[i + j], x = a[i + j + half];
                u32 v = mulm(x, (u32)cur);
                a[i + j] = addm(u, v); a[i + j + half] = subm(u, v);
                cur = cur * lw % MOD;
            }
        }
    }
    u32 ninv = fpow((u32)n, MOD - 2);
    for (int i = 0; i < n; i++) a[i] = mulm(a[i], ninv);
}

AVX2FN static void tail8fwd(u32 *a, int n) {
    u32 g8 = fpow(3, (MOD - 1) / 8), g4 = fpow(3, (MOD - 1) / 4);
    __m256i W8[4], S8[4], W4[2], S4[2];
    for (int j = 0; j < 4; j++) { u32 w = fpow(g8, (u32)j); W8[j] = _mm256_set1_epi32((int)w);
        S8[j] = _mm256_set1_epi32((int)(u32)(((__uint128_t)w << 32) / MOD)); }
    for (int j = 0; j < 2; j++) { u32 w = fpow(g4, (u32)j); W4[j] = _mm256_set1_epi32((int)w);
        S4[j] = _mm256_set1_epi32((int)(u32)(((__uint128_t)w << 32) / MOD)); }
    u32 w8c[4], w4c[2];
    for (int j = 0; j < 4; j++) w8c[j] = fpow(g8, (u32)j);
    for (int j = 0; j < 2; j++) w4c[j] = fpow(g4, (u32)j);
    for (int i = 0; i + 64 <= n; i += 64) {
        __m256i v[8];
        for (int m = 0; m < 8; m++) v[m] = _mm256_loadu_si256((const __m256i*)(const void*)(a + i + 8 * m));
        tr8i(v);
        for (int j = 0; j < 4; j++) { __m256i u = v[j], x = v[j + 4]; v[j] = addmod8(u, x); v[j + 4] = shoup8(submod8(u, x), W8[j], S8[j]); }
        for (int b = 0; b < 8; b += 4) for (int j = 0; j < 2; j++) { __m256i u = v[b + j], x = v[b + j + 2]; v[b + j] = addmod8(u, x); v[b + j + 2] = shoup8(submod8(u, x), W4[j], S4[j]); }
        for (int b = 0; b < 8; b += 2) { __m256i u = v[b], x = v[b + 1]; v[b] = addmod8(u, x); v[b + 1] = submod8(u, x); }
        tr8i(v);
        for (int m = 0; m < 8; m++) _mm256_storeu_si256((__m256i*)(void*)(a + i + 8 * m), v[m]);
    }
    for (int i = (n / 64) * 64; i + 8 <= n; i += 8) k8fwd_scalar(a + i, w8c, w4c);
}
AVX2FN static void tail8inv(u32 *a, int n) {
    u32 g8 = fpow(3, MOD - 1 - (MOD - 1) / 8), g4 = fpow(3, MOD - 1 - (MOD - 1) / 4);
    __m256i W8[4], S8[4], W4[2], S4[2];
    for (int j = 0; j < 4; j++) { u32 w = fpow(g8, (u32)j); W8[j] = _mm256_set1_epi32((int)w);
        S8[j] = _mm256_set1_epi32((int)(u32)(((__uint128_t)w << 32) / MOD)); }
    for (int j = 0; j < 2; j++) { u32 w = fpow(g4, (u32)j); W4[j] = _mm256_set1_epi32((int)w);
        S4[j] = _mm256_set1_epi32((int)(u32)(((__uint128_t)w << 32) / MOD)); }
    for (int i = 0; i + 64 <= n; i += 64) {
        __m256i v[8];
        for (int m = 0; m < 8; m++) v[m] = _mm256_loadu_si256((const __m256i*)(const void*)(a + i + 8 * m));
        tr8i(v);
        for (int b = 0; b < 8; b += 2) { __m256i u = v[b], x = v[b + 1]; v[b] = addmod8(u, x); v[b + 1] = submod8(u, x); }
        for (int b = 0; b < 8; b += 4) for (int j = 0; j < 2; j++) { __m256i u = v[b + j], x = shoup8(v[b + j + 2], W4[j], S4[j]); v[b + j] = addmod8(u, x); v[b + j + 2] = submod8(u, x); }
        for (int j = 0; j < 4; j++) { __m256i u = v[j], x = shoup8(v[j + 4], W8[j], S8[j]); v[j] = addmod8(u, x); v[j + 4] = submod8(u, x); }
        tr8i(v);
        for (int m = 0; m < 8; m++) _mm256_storeu_si256((__m256i*)(void*)(a + i + 8 * m), v[m]);
    }
    u32 w8i[4], w4i[2];
    for (int j = 0; j < 4; j++) w8i[j] = fpow(g8, (u32)j);
    for (int j = 0; j < 2; j++) w4i[j] = fpow(g4, (u32)j);
    for (int i = (n / 64) * 64; i + 8 <= n; i += 8) k8inv_scalar(a + i, w8i, w4i);
}

AVX2FN static void ntt2_fwd(u32 *a, int n, const Tables2 &T) {
    if (n < 16) { small_fwd(a, n); return; }
    for (int len = n; len >= 16; len >>= 1) {
        int half = len >> 1;
        const u32 *w = T.fw + T.foff[__builtin_ctz(len)];
        const u32 *ws = T.fws + T.foff[__builtin_ctz(len)];
        for (int i = 0; i < n; i += len) {
            u32 *lo = a + i, *hi = lo + half;
            int j = 0;
            for (; j + 8 <= half; j += 8) {
                __m256i u = _mm256_loadu_si256((const __m256i*)(const void*)(lo + j));
                __m256i v = _mm256_loadu_si256((const __m256i*)(const void*)(hi + j));
                __m256i W = _mm256_loadu_si256((const __m256i*)(const void*)(w + j));
                _mm256_storeu_si256((__m256i*)(void*)(lo + j), addmod8(u, v));
                _mm256_storeu_si256((__m256i*)(void*)(hi + j), shoup8(submod8(u, v), W, _mm256_loadu_si256((const __m256i*)(const void*)(ws + j))));
            }
            for (; j < half; j++) {
                u32 u = lo[j], v = hi[j];
                u32 wsj = (u32)(((__uint128_t)w[j] << 32) / MOD);
                u32 q = (u32)(((u64)subm(u, v) * wsj) >> 32);
                u32 r = (u32)((u64)subm(u, v) * w[j] - (u64)q * MOD);
                if (r >= MOD) r -= MOD;
                lo[j] = addm(u, v); hi[j] = r;
            }
        }
    }
    tail8fwd(a, n);
}
AVX2FN static void ntt2_inv(u32 *a, int n, const Tables2 &T) {
    if (n < 16) { small_inv(a, n); return; }
    u32 ninv = fpow((u32)n, MOD - 2);
    u32 ninvs = (u32)(((__uint128_t)ninv << 32) / MOD);
    __m256i NINV = _mm256_set1_epi32((int)ninv), NINVS = _mm256_set1_epi32((int)ninvs);
    tail8inv(a, n);
    for (int len = 16; len <= n; len <<= 1) {
        int half = len >> 1;
        const u32 *w = T.iw + T.ioff[__builtin_ctz(len)];
        const u32 *ws = T.iws + T.ioff[__builtin_ctz(len)];
        for (int i = 0; i < n; i += len) {
            u32 *lo = a + i, *hi = lo + half;
            int j = 0;
            for (; j + 8 <= half; j += 8) {
                __m256i u = _mm256_loadu_si256((const __m256i*)(const void*)(lo + j));
                __m256i x = _mm256_loadu_si256((const __m256i*)(const void*)(hi + j));
                __m256i W = _mm256_loadu_si256((const __m256i*)(const void*)(w + j));
                if (len == n) { u = shoup8(u, NINV, NINVS); x = shoup8(x, NINV, NINVS); }
                __m256i v = shoup8(x, W, _mm256_loadu_si256((const __m256i*)(const void*)(ws + j)));
                _mm256_storeu_si256((__m256i*)(void*)(lo + j), addmod8(u, v));
                _mm256_storeu_si256((__m256i*)(void*)(hi + j), submod8(u, v));
            }
            for (; j < half; j++) {
                u32 u = lo[j], x = hi[j];
                if (len == n) { u = mulm(u, ninv); x = mulm(x, ninv); }
                u32 wsj = (u32)(((__uint128_t)w[j] << 32) / MOD);
                u32 q = (u32)(((u64)x * wsj) >> 32);
                u32 v = (u32)((u64)x * w[j] - (u64)q * MOD);
                if (v >= MOD) v -= MOD;
                lo[j] = addm(u, v); hi[j] = subm(u, v);
            }
        }
    }
}



/* ---------- 1002: four-step (cache-blocked) NTT for N = N1*N2 ---------- */
static u32 *gBuf, *gScratch, *gOut, *gSpec, *gTw;
static Tables2 T1, T2;
static int GN1, GN2, GN;

static u32 *bigAlloc(size_t sz) {
    void *p = 0;
    if (posix_memalign(&p, 64, sz)) return 0;
    memset(p, 0, sz);
    return (u32 *)p;
}

static void transpose_blk(const u32 *src, u32 *dst, int R, int C) {
    const int B = 32;
    u32 tmp[32][32];
    for (int i0 = 0; i0 < R; i0 += B)
        for (int j0 = 0; j0 < C; j0 += B) {
            int ie = i0 + B < R ? i0 + B : R, je = j0 + B < C ? j0 + B : C;
            /* read the block rows (contiguous 128B each) */
            for (int i = i0; i < ie; i++) memcpy(tmp[i - i0], src + (size_t)i * C + j0, (size_t)(je - j0) * 4);
            /* write the transposed rows (contiguous) */
            for (int j = j0; j < je; j++) {
                u32 *d = dst + (size_t)j * R + i0;
                int k = j - j0;
                for (int i = 0; i < ie - i0; i++) d[i] = tmp[i][k];
            }
        }
}

static inline int bitrev(int v, int bits) {
    int r = 0;
    for (int b = 0; b < bits; b++) { r = (r << 1) | (v & 1); v >>= 1; }
    return r;
}

/* forward DFT, natural order in -> natural order out (permutation applied at the end) */
AVX2FN static void fourstep_fwd(u32 *a) {
    /* step 1: B'[i2][i1] = a[i1][i2] */
    transpose_blk(a, gScratch, GN1, GN2);            /* (N1 x N2) -> (N2 x N1) */
    /* step 2: length-N1 DFT on each row of B' */
    for (int i2 = 0; i2 < GN2; i2++) ntt2_fwd(gScratch + (size_t)i2 * GN1, GN1, T1);
    /* step 3: twiddle by w_N^{i2 * k1}, k1 = bitrev(j1) */
    {
        const u32 *tw = gTw;
        for (int i2 = 0; i2 < GN2; i2++) {
            u32 *row = gScratch + (size_t)i2 * GN1;
            const u32 *t = tw + (size_t)i2 * GN1;
            for (int j = 0; j < GN1; j += 8) {
                __m256i v = _mm256_loadu_si256((const __m256i *)(const void *)(row + j));
                __m256i W = _mm256_loadu_si256((const __m256i *)(const void *)(t + j));
                __m256i WS = _mm256_loadu_si256((const __m256i *)(const void *)(t + GN + j));
                _mm256_storeu_si256((__m256i *)(void *)(row + j), shoup8(v, W, WS));
            }
        }
    }
    /* step 4: transpose back so rows are indexed by j1 with contiguous i2 */
    transpose_blk(gScratch, a, GN2, GN1);            /* (N2 x N1) -> (N1 x N2) */
    /* step 5: length-N2 DFT on each row */
    for (int j1 = 0; j1 < GN1; j1++) ntt2_fwd(a + (size_t)j1 * GN2, GN2, T2);
}

/* inverse of the pipeline above: consumes the P-ordered spectrum and returns
   the time domain, so no bit-reversal permutation pass is needed anywhere. */
AVX2FN static void fourstep_inv(u32 *a) {
    for (int j1 = 0; j1 < GN1; j1++) ntt2_inv(a + (size_t)j1 * GN2, GN2, T2);
    transpose_blk(a, gScratch, GN1, GN2);            /* [j1][i2] -> [i2][j1] */
    {
        const u32 *tw = gTw + (size_t)GN * 2;        /* inverse twiddles live after the forward pair */
        const u32 *tws = tw + GN;
        for (int i2 = 0; i2 < GN2; i2++) {
            u32 *row = gScratch + (size_t)i2 * GN1;
            const u32 *t = tw + (size_t)i2 * GN1;
            const u32 *ts = tws + (size_t)i2 * GN1;
            for (int j = 0; j < GN1; j += 8) {
                __m256i v = _mm256_loadu_si256((const __m256i *)(const void *)(row + j));
                __m256i W = _mm256_loadu_si256((const __m256i *)(const void *)(t + j));
                __m256i WS = _mm256_loadu_si256((const __m256i *)(const void *)(ts + j));
                _mm256_storeu_si256((__m256i *)(void *)(row + j), shoup8(v, W, WS));
            }
        }
    }
    for (int i2 = 0; i2 < GN2; i2++) ntt2_inv(gScratch + (size_t)i2 * GN1, GN1, T1);
    transpose_blk(gScratch, a, GN2, GN1);            /* [i2][i1] -> [i1][i2] */
    /* note: ntt2_inv already carries 1/N1 and 1/N2, so the result is scaled by 1/N */
}

static void fourstep_prepare(int n1, int n2) {
    GN1 = n1; GN2 = n2; GN = n1 * n2;
    T1.init(n1); T2.init(n2);
    gBuf = bigAlloc((size_t)GN * 4);
    gScratch = bigAlloc((size_t)GN * 4);
    gOut = bigAlloc((size_t)GN * 4);
    gSpec = bigAlloc((size_t)GN * 4);
    gTw = bigAlloc((size_t)GN * 16);                 /* [fwd | fwds | inv | invs] */
    /* tw[i2*N1 + j1] = w_N^{i2 * bitrev(j1)} ; tws = Shoup constant */
    u32 wN = fpow(3, (MOD - 1) / GN);
    u32 wNi = fpow(wN, MOD - 2);
    const double inv = 4294967296.0 / (double)MOD;
    /* write row by row (contiguous) and only scatter *inside* the 8KB row */
    for (int i2 = 0; i2 < GN2; i2++) {
        u32 base = fpow(wN, (u32)i2), basei = fpow(wNi, (u32)i2);
        u64 cur = 1, curi = 1;
        u32 *rowf = gTw + (size_t)i2 * GN1;
        u32 *rows = gTw + (size_t)GN + (size_t)i2 * GN1;
        u32 *rowfi = gTw + (size_t)GN * 2 + (size_t)i2 * GN1;
        u32 *rowsi = gTw + (size_t)GN * 3 + (size_t)i2 * GN1;
        for (int k1 = 0; k1 < GN1; k1++) {
            int j1 = bitrev(k1, POLY_N1BITS);
            u32 x = (u32)cur;
            u32 q = (u32)((double)x * inv);
            u64 lhs = (u64)q * MOD, rhs = (u64)x << 32;
            if (lhs > rhs) q--; else if (lhs + MOD <= rhs) q++;
            rowf[j1] = x; rows[j1] = q;
            u32 xi = (u32)curi;
            u32 qi = (u32)((double)xi * inv);
            u64 lhsi = (u64)qi * MOD, rhsi = (u64)xi << 32;
            if (lhsi > rhsi) qi--; else if (lhsi + MOD <= rhsi) qi++;
            rowfi[j1] = xi; rowsi[j1] = qi;
            cur = cur * base % MOD;
            curi = curi * basei % MOD;
        }
    }
}

extern "C" void poly_multiply(unsigned *a, int n, unsigned *b, int m, unsigned *c) {
    const int N1 = 1 << POLY_N1BITS, N2 = 1 << POLY_N2BITS, N = N1 * N2;
    fourstep_prepare(N1, N2);
    if (!gBuf || !gScratch || !gSpec || !gTw) return;
    int la = (n + 1 < N) ? n + 1 : N;
    int lb = (m + 1 < N) ? m + 1 : N;
    for (int i = 0; i < N; i++) gBuf[i] = 0;
    for (int i = 0; i < la; i++) gBuf[i] = a[i];
    fourstep_fwd(gBuf);                     /* P-ordered spectrum of A */
    memcpy(gSpec, gBuf, (size_t)N * 4);
    for (int i = 0; i < N; i++) gBuf[i] = 0;
    for (int i = 0; i < lb; i++) gBuf[i] = b[i];
    fourstep_fwd(gBuf);                     /* P-ordered spectrum of B */
    for (int i = 0; i < N; i++) gBuf[i] = mulm(gSpec[i], gBuf[i]);
    fourstep_inv(gBuf);                     /* back to the time domain, already scaled */
    int outc = n + m + 1;
    if (outc > N) outc = N;
    for (int i = 0; i < outc; i++) c[i] = gBuf[i];
}

CompilationN/AN/ACompile ErrorScore: N/A


Judge Duck Online | 评测鸭在线
Server Time: 2026-11-06 15:23:17 | Loaded in 1 ms | Server Status
个人娱乐项目,仅供学习交流使用 | 捐赠