代码

// 1004: 4-lane single-mod AVX2 NTT (Montgomery), 3 mods, base 1e9, NTT 2^18.
// DIF forward + DIT inverse. Big stages (len>=8) vectorized 4-lane; small stages (len=4,2) scalar.
#include <sys/auxv.h>
#include <stdint.h>
#include <string.h>
#include <immintrin.h>
#pragma GCC target("avx2")

typedef uint64_t u64;
typedef uint32_t u32;
typedef __uint128_t u128;

struct DuckInfo {
  uint64_t abi_version;
  const char *stdin_ptr; uint64_t stdin_size;
  char *stdout_ptr; uint64_t stdout_limit; uint64_t stdout_size;
  char *stderr_ptr; uint64_t stderr_limit; uint64_t stderr_size;
  const char *IB_ptr; uint64_t IB_limit;
  char *OB_ptr; uint64_t OB_limit;
  uint64_t tsc_frequency;
} __attribute__((packed));

static const u32 MODS[3] = {998244353u, 1004535809u, 469762049u};
static const u32 NINV[3] = {998244351u, 1004535807u, 469762047u};
static const u32 R2[3]   = {932051910u, 542374313u, 460175152u};

#define SIZE (1<<18)
#define VSIZE (SIZE/4)

static __m256i A[VSIZE] __attribute__((aligned(64)));
static __m256i B[VSIZE] __attribute__((aligned(64)));
static u32 AU[SIZE], BU[SIZE];
static u32 TW[SIZE], ITW[SIZE];
static u32 limbsA[111112], limbsB[111112];
static u32 outlimbs[222224];
static u32 r0[SIZE], r1[SIZE], r2[SIZE];

static __m256i modvec, ninvvec, sub1vec, MASK32, ZERO;

static u32 powmod(u64 a, u64 e, u32 mod) {
  u64 r = 1, b = a % mod;
  while (e) { if (e & 1) r = r * b % mod; b = b * b % mod; e >>= 1; }
  return (u32)r;
}
static inline u32 mont_s(u32 a, u32 b, u32 mod, u32 ninv) {
  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 u32 sadd(u32 a, u32 b, u32 mod) { u32 s = a + b; if (s >= mod) s -= mod; return s; }
static inline u32 ssub(u32 a, u32 b, u32 mod) { u32 d = a - b; if (d > mod) d += mod; return d; }

// fill dst[0..half) = wstep^j (Montgomery), block-parallel to avoid serial dependency
static void fill_stage(u32* dst, int half, u32 wstep, u32 onem, u32 mod, u32 ninv) {
  const int B = 16;
  dst[0] = onem;
  int k = 1;
  for (; k < B && k < half; k++) dst[k] = mont_s(dst[k - 1], wstep, mod, ninv);
  if (half <= B) return;
  u32 wB = wstep;
  for (int t = 1; t < B; t++) wB = mont_s(wB, wstep, mod, ninv);
  for (int base = B; base < half; base += B) {
    u32 wbase = mont_s(dst[base - B], wB, mod, ninv);
    dst[base] = wbase;
    int lim = (base + B < half) ? B : (half - base);
    for (int j = 1; j < lim; j++) dst[base + j] = mont_s(wbase, dst[j], mod, ninv);
  }
}

static inline __m256i mont_mul(__m256i a, __m256i b) {
  __m256i t = _mm256_mul_epu32(a, b);
  __m256i m = _mm256_mul_epu32(t, ninvvec);
  __m256i u = _mm256_srli_epi64(_mm256_add_epi64(t, _mm256_mul_epu32(m, modvec)), 32);
  __m256i ge = _mm256_cmpgt_epi32(u, sub1vec);
  return _mm256_sub_epi32(u, _mm256_and_si256(ge, modvec));
}
static inline __m256i vadd(__m256i a, __m256i b) {
  __m256i s = _mm256_add_epi32(a, b);
  __m256i ge = _mm256_cmpgt_epi32(s, sub1vec);
  return _mm256_sub_epi32(s, _mm256_and_si256(ge, modvec));
}
static inline __m256i vsub(__m256i a, __m256i b) {
  __m256i d = _mm256_sub_epi32(a, b);
  __m256i lt = _mm256_cmpgt_epi32(ZERO, d);
  return _mm256_add_epi32(d, _mm256_and_si256(lt, modvec));
}

static inline __m256i load_tw(const u32* tw, int j) {
  __m128i lo = _mm_loadu_si128((const __m128i*)(tw + j));
  return _mm256_cvtepu32_epi64(lo);
}

static void vec_to_u32(const __m256i* src, u32* dst) {
  __m256i perm = _mm256_setr_epi32(0, 2, 4, 6, 0, 0, 0, 0);
  for (int v = 0; v < VSIZE; v++) {
    __m256i t = _mm256_permutevar8x32_epi32(src[v], perm);
    _mm_storeu_si128((__m128i*)(dst + v * 4), _mm256_castsi256_si128(t));
  }
}
static void u32_to_vec(const u32* src, __m256i* dst) {
  for (int v = 0; v < VSIZE; v++) {
    __m128i lo = _mm_loadu_si128((const __m128i*)(src + v * 4));
    dst[v] = _mm256_cvtepu32_epi64(lo);
  }
}

// DIF forward big stages (len = SIZE .. 8). Returns twiddle offset after big stages.
static int ntt_fwd_big(__m256i* x, const u32* TW) {
  int n = SIZE, off = 0;
  for (int len = n; len >= 8; len >>= 1) {
    int half = len >> 1;
    int hv = half >> 2; // half in vector units
    for (int i = 0; i < n; i += len) {
      __m256i* y = x + (i >> 2);
      int b = 0;
      for (; b + 4 <= hv; b += 4) {
        __m256i u0 = y[b],       v0 = y[b + hv];
        __m256i u1 = y[b + 1],   v1 = y[b + 1 + hv];
        __m256i u2 = y[b + 2],   v2 = y[b + 2 + hv];
        __m256i u3 = y[b + 3],   v3 = y[b + 3 + hv];
        __m256i tw0 = load_tw(TW, off + (b << 2));
        __m256i tw1 = load_tw(TW, off + (b << 2) + 4);
        __m256i tw2 = load_tw(TW, off + (b << 2) + 8);
        __m256i tw3 = load_tw(TW, off + (b << 2) + 12);
        y[b]       = vadd(u0, v0);  y[b + hv]       = mont_mul(vsub(u0, v0), tw0);
        y[b + 1]   = vadd(u1, v1);  y[b + 1 + hv]   = mont_mul(vsub(u1, v1), tw1);
        y[b + 2]   = vadd(u2, v2);  y[b + 2 + hv]   = mont_mul(vsub(u2, v2), tw2);
        y[b + 3]   = vadd(u3, v3);  y[b + 3 + hv]   = mont_mul(vsub(u3, v3), tw3);
      }
      for (; b < hv; b++) {
        __m256i u = y[b], v = y[b + hv], tw = load_tw(TW, off + (b << 2));
        y[b] = vadd(u, v);
        y[b + hv] = mont_mul(vsub(u, v), tw);
      }
    }
    off += half;
  }
  return off;
}

// DIF forward small stages (len=4, len=2) on u32 array. off is twiddle offset.
static void ntt_fwd_small(u32* x, const u32* TW, int off, u32 mod, u32 ninv) {
  // len=4
  {
    int step = SIZE / 4;
    u32 tw1 = TW[off + 1]; // w^step (Montgomery); tw0 = TW[off] = 1
    for (int i = 0; i < SIZE; i += 4) {
      u32 u0 = x[i], v0 = x[i + 2];
      u32 u1 = x[i + 1], v1 = x[i + 3];
      x[i] = sadd(u0, v0, mod);
      x[i + 2] = mont_s(ssub(u0, v0, mod), TW[off], mod, ninv);
      x[i + 1] = sadd(u1, v1, mod);
      x[i + 3] = mont_s(ssub(u1, v1, mod), tw1, mod, ninv);
    }
    off += 2;
  }
  // len=2
  {
    for (int i = 0; i < SIZE; i += 2) {
      u32 u0 = x[i], v0 = x[i + 1];
      x[i] = sadd(u0, v0, mod);
      x[i + 1] = ssub(u0, v0, mod);
    }
    off += 1;
  }
  (void)off;
}

// DIT inverse small stages (len=2, len=4) on u32 array. ITW starts at off=0.
static void ntt_inv_small(u32* x, const u32* ITW, u32 mod, u32 ninv) {
  int off = 0;
  // len=2
  {
    for (int i = 0; i < SIZE; i += 2) {
      u32 u0 = x[i], v0 = x[i + 1];
      x[i] = sadd(u0, v0, mod);
      x[i + 1] = ssub(u0, v0, mod);
    }
    off += 1;
  }
  // len=4
  {
    u32 tw1 = ITW[off + 1];
    for (int i = 0; i < SIZE; i += 4) {
      u32 u0 = x[i], v0 = mont_s(x[i + 2], ITW[off], mod, ninv);
      u32 u1 = x[i + 1], v1 = mont_s(x[i + 3], tw1, mod, ninv);
      x[i] = sadd(u0, v0, mod);
      x[i + 2] = ssub(u0, v0, mod);
      x[i + 1] = sadd(u1, v1, mod);
      x[i + 3] = ssub(u1, v1, mod);
    }
    off += 2;
  }
}

// DIT inverse big stages (len=8 .. SIZE). ITW offset starts after small stages (=3).
static void ntt_inv_big(__m256i* x, const u32* ITW) {
  int n = SIZE, off = 3;
  for (int len = 8; len <= n; len <<= 1) {
    int half = len >> 1;
    int hv = half >> 2;
    for (int i = 0; i < n; i += len) {
      __m256i* y = x + (i >> 2);
      int b = 0;
      for (; b + 4 <= hv; b += 4) {
        __m256i u0 = y[b],       v0 = mont_mul(y[b + hv],       load_tw(ITW, off + (b << 2)));
        __m256i u1 = y[b + 1],   v1 = mont_mul(y[b + 1 + hv],   load_tw(ITW, off + (b << 2) + 4));
        __m256i u2 = y[b + 2],   v2 = mont_mul(y[b + 2 + hv],   load_tw(ITW, off + (b << 2) + 8));
        __m256i u3 = y[b + 3],   v3 = mont_mul(y[b + 3 + hv],   load_tw(ITW, off + (b << 2) + 12));
        y[b]       = vadd(u0, v0);  y[b + hv]       = vsub(u0, v0);
        y[b + 1]   = vadd(u1, v1);  y[b + 1 + hv]   = vsub(u1, v1);
        y[b + 2]   = vadd(u2, v2);  y[b + 2 + hv]   = vsub(u2, v2);
        y[b + 3]   = vadd(u3, v3);  y[b + 3 + hv]   = vsub(u3, v3);
      }
      for (; b < hv; b++) {
        __m256i u = y[b], v = mont_mul(y[b + hv], load_tw(ITW, off + (b << 2)));
        y[b] = vadd(u, v);
        y[b + hv] = vsub(u, v);
      }
    }
    off += half;
  }
}

static char tab3[1000][3];
static void build_tab3(void) {
  for (int i = 0; i < 1000; i++) {
    int v = i;
    tab3[i][2] = '0' + v % 10; v /= 10;
    tab3[i][1] = '0' + v % 10; v /= 10;
    tab3[i][0] = '0' + v % 10;
  }
}

#ifdef LOCAL_TEST
extern uintptr_t jd_getauxval(uintptr_t);
int jd_main() {
  DuckInfo* di = (DuckInfo*)jd_getauxval(0x6b637564ull);
#else
int main() {
  DuckInfo* di = (DuckInfo*)getauxval(0x6b637564ull);
#endif
  const char* in = di->stdin_ptr;
  u64 inlen = di->stdin_size;

  const char* p = in;
  const char* inend = in + inlen;
  while (p < inend && (*p == '\n' || *p == '\r' || *p == ' ' || *p == '\t')) p++;
  const char* a_start = p;
  while (p < inend && *p >= '0' && *p <= '9') p++;
  const char* a_end = p;
  while (p < inend && (*p == '\n' || *p == '\r' || *p == ' ' || *p == '\t')) p++;
  const char* b_start = p;
  while (p < inend && *p >= '0' && *p <= '9') p++;
  const char* b_end = p;

  const char* sa = a_start;
  while (sa < a_end - 1 && *sa == '0') sa++;
  const char* sb = b_start;
  while (sb < b_end - 1 && *sb == '0') sb++;

  int na = 0, nb = 0;
  {
    const char* pos = a_end;
    while (pos > sa) {
      const char* st = pos - 9; if (st < sa) st = sa;
      u32 v = 0;
      for (const char* q = st; q < pos; q++) v = v * 10 + (u32)(*q - '0');
      limbsA[na++] = v;
      pos = st;
    }
  }
  {
    const char* pos = b_end;
    while (pos > sb) {
      const char* st = pos - 9; if (st < sb) st = sb;
      u32 v = 0;
      for (const char* q = st; q < pos; q++) v = v * 10 + (u32)(*q - '0');
      limbsB[nb++] = v;
      pos = st;
    }
  }

  MASK32 = _mm256_set1_epi64x(0xFFFFFFFFLL);
  ZERO = _mm256_setzero_si256();

  for (int mi = 0; mi < 3; mi++) {
    u32 mod = MODS[mi], ninv = NINV[mi], r2c = R2[mi];
    modvec = _mm256_set1_epi64x((long long)mod);
    ninvvec = _mm256_set1_epi64x((long long)ninv);
    sub1vec = _mm256_set1_epi64x((long long)(mod - 1));

    u32 w = powmod(3u, (mod - 1) / SIZE, mod);
    u32 iw = powmod(w, mod - 2, mod);
    u32 onem = mont_s(1, r2c, mod, ninv); // 1 in Montgomery form

    // build TW (DIF forward order, len = SIZE..2)
    {
      int off = 0;
      for (int len = SIZE; len > 1; len >>= 1) {
        int half = len >> 1;
        int step = SIZE / len;
        u32 wstep = mont_s(powmod(w, step, mod), r2c, mod, ninv);
        fill_stage(TW + off, half, wstep, onem, mod, ninv);
        off += half;
      }
    }
    // build ITW (DIT inverse order, len = 2..SIZE)
    {
      int off = 0;
      for (int len = 2; len <= SIZE; len <<= 1) {
        int half = len >> 1;
        int step = SIZE / len;
        u32 wstep = mont_s(powmod(iw, step, mod), r2c, mod, ninv);
        fill_stage(ITW + off, half, wstep, onem, mod, ninv);
        off += half;
      }
    }

    // fill A, B (vector, Montgomery form)
    __m256i r2v = _mm256_set1_epi64x((long long)r2c);
    {
      int vlim = (na + 3) >> 2;
      for (int v = 0; v < vlim; v++) {
        u32 x[4];
        for (int k = 0; k < 4; k++) {
          int idx = v * 4 + k;
          u32 val = (idx < na) ? limbsA[idx] : 0u;
          while (val >= mod) val -= mod;
          x[k] = val;
        }
        A[v] = mont_mul(_mm256_setr_epi64x((long long)x[0], (long long)x[1], (long long)x[2], (long long)x[3]), r2v);
      }
      for (int v = vlim; v < VSIZE; v++) A[v] = ZERO;
    }
    {
      int vlim = (nb + 3) >> 2;
      for (int v = 0; v < vlim; v++) {
        u32 x[4];
        for (int k = 0; k < 4; k++) {
          int idx = v * 4 + k;
          u32 val = (idx < nb) ? limbsB[idx] : 0u;
          while (val >= mod) val -= mod;
          x[k] = val;
        }
        B[v] = mont_mul(_mm256_setr_epi64x((long long)x[0], (long long)x[1], (long long)x[2], (long long)x[3]), r2v);
      }
      for (int v = vlim; v < VSIZE; v++) B[v] = ZERO;
    }

    // forward
    int off = ntt_fwd_big(A, TW);
    ntt_fwd_big(B, TW);
    vec_to_u32(A, AU);
    vec_to_u32(B, BU);
    ntt_fwd_small(AU, TW, off, mod, ninv);
    ntt_fwd_small(BU, TW, off, mod, ninv);

    // pointwise (scalar, Montgomery)
    for (int i = 0; i < SIZE; i++) AU[i] = mont_s(AU[i], BU[i], mod, ninv);

    // inverse
    ntt_inv_small(AU, ITW, mod, ninv);
    u32_to_vec(AU, A);
    ntt_inv_big(A, ITW);

    // scale by n^-1, extract residues
    u32 ninvn = powmod(SIZE, mod - 2, mod);
    __m256i ninvnv = _mm256_set1_epi64x((long long)ninvn);
    for (int v = 0; v < VSIZE; v++) A[v] = mont_mul(A[v], ninvnv);
    u32* dst = (mi == 0) ? r0 : (mi == 1) ? r1 : r2;
    vec_to_u32(A, dst);
  }

  // CRT reconstruction (Garner), carry in base 1e9, then output
  {
    u32 m0 = MODS[0], m1 = MODS[1], m2 = MODS[2];
    u64 M0 = m0, M1 = m1;
    u64 INV01 = powmod(M0 % m1, m1 - 2, m1);
    u128 M01 = (u128)M0 * M1;
    u64 INV012 = powmod((u64)(M01 % m2), m2 - 2, m2);

    int outlen = na + nb - 1;
    u128 carry = 0;
    for (int i = 0; i < outlen; i++) {
      u32 a0 = r0[i], a1 = r1[i], a2 = r2[i];
      u128 c = a0;
      u64 t1 = (u64)(a1 - (u32)(c % m1) + m1) % m1;
      t1 = (u64)((u128)t1 * INV01 % m1);
      c += (u128)t1 * M0;
      u64 t2 = (u64)(a2 - (u32)(c % m2) + m2) % m2;
      t2 = (u64)((u128)t2 * INV012 % m2);
      c += (u128)t2 * M01;
      c += carry;
      outlimbs[i] = (u32)(c % 1000000000u);
      carry = c / 1000000000u;
    }
    while (carry) { outlimbs[outlen++] = (u32)(carry % 1000000000u); carry /= 1000000000u; }

    int hi = outlen - 1;
    while (hi > 0 && outlimbs[hi] == 0) hi--;

    build_tab3();
    char* out = di->stdout_ptr;
    char* o = out;
    {
      u32 v = outlimbs[hi];
      char tmp[10]; int t = 0;
      do { tmp[t++] = '0' + (char)(v % 10); v /= 10; } while (v);
      while (t > 0) *o++ = tmp[--t];
    }
    for (int i = hi - 1; i >= 0; i--) {
      u32 v = outlimbs[i];
      u32 g2 = v / 1000000u;
      u32 g1 = (v / 1000u) % 1000u;
      u32 g0 = v % 1000u;
      const char* q;
      q = tab3[g2]; *o++ = q[0]; *o++ = q[1]; *o++ = q[2];
      q = tab3[g1]; *o++ = q[0]; *o++ = q[1]; *o++ = q[2];
      q = tab3[g0]; *o++ = q[0]; *o++ = q[1]; *o++ = q[2];
    }
    *o++ = '\n';
    di->stdout_size = (u64)(o - out);
  }

#ifdef LOCAL_TEST
  return 0;
#else
  asm volatile("mov $60, %%eax; xor %%edi, %%edi; syscall" ::: "rax", "rdi", "rcx", "r11", "memory");
  __builtin_unreachable();
#endif
}

评测结果