// 1004: radix-4 (2 x 4^7 x 8) DIF/DIT, 3-mod NTT (998244353, 1004535809, 469762049),
// all-Montgomery, interleaved twiddles, base 1e9 limbs, DuckInfo direct IO.
// Structure: 1 radix-2 DIF + 7 radix-4 DIF + ntt8_vec base (forward),
// ntt8_vec base + 7 radix-4 DIT + 1 radix-2 DIT (inverse).
#include <sys/auxv.h>
#include <stdint.h>
#include <string.h>
#include <immintrin.h>
#pragma GCC target("avx2")
#pragma GCC optimize("O3,unroll-loops")
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 NINVS[3] = {998244351u, 1004535807u, 469762047u};
static const u32 R2S[3] = {932051910u, 542374313u, 460175152u};
static const u32 ZETA[3] = {372528824u, 395918948u, 129701348u}; // primitive 8th root (std)
static const u32 ZETA3[3] = {488723995u, 823254837u, 443138433u};
static const u32 ZETA5[3] = {625715529u, 608616861u, 340060701u};
static const u32 ZETA7[3] = {509520358u, 181280972u, 26623616u};
static const u32 IROOT[3] = {911660635u, 483363861u, 450151958u}; // sqrt(-1) (std)
static const u32 IMINUS[3]= {86583718u, 521171948u, 19610091u};
#define SIZE (1<<18)
static u32 A[SIZE], B[SIZE];
static u32 TW[SIZE+64], ITW[SIZE+64];
static u32 tmp[SIZE/4+8];
static u32 limbsA[111112], limbsB[111112];
static u32 outlimbs[222224];
static u32 r0[SIZE], r1[SIZE], r2[SIZE];
static char tab3[1000][3];
__m256i modvec, ninvvec, shufvec;
__m256i z1v, z3v, ivv;
static u32 curmod, curninv;
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; }
static inline __m256i mont8(__m256i x, __m256i y) {
__m256i ao=_mm256_srli_epi64(x,32);
__m256i bo=_mm256_srli_epi64(y,32);
__m256i te=_mm256_mul_epu32(x,y);
__m256i to=_mm256_mul_epu32(ao,bo);
__m256i me=_mm256_mul_epu32(te,ninvvec);
__m256i mo=_mm256_mul_epu32(to,ninvvec);
__m256i ue=_mm256_srli_epi64(_mm256_add_epi64(te,_mm256_mul_epu32(me,modvec)),32);
__m256i uo=_mm256_srli_epi64(_mm256_add_epi64(to,_mm256_mul_epu32(mo,modvec)),32);
__m256i u=_mm256_or_si256(ue,_mm256_slli_epi64(uo,32));
return _mm256_min_epu32(u,_mm256_sub_epi32(u,modvec));
}
static inline __m256i addm(__m256i a, __m256i b){ __m256i s=_mm256_add_epi32(a,b); return _mm256_min_epu32(s,_mm256_sub_epi32(s,modvec)); }
static inline __m256i subm(__m256i a, __m256i b){ __m256i d=_mm256_sub_epi32(a,b); return _mm256_min_epu32(d,_mm256_add_epi32(d,modvec)); }
static inline __m256i loadu(const u32* p){ return _mm256_loadu_si256((const __m256i*)p); }
static inline void storeu(u32* p, __m256i v){ _mm256_storeu_si256((__m256i*)p, v); }
// Vectorized geometric progression: tmp[p] = mont(wm^p), p=0..s-1 (mont8 globals must be set).
static void gen_powers(u32* tmp, int s, u32 wm, u32 mod, u32 ninv, u32 r2c) {
tmp[0]=mont_s(1,r2c,mod,ninv);
u32 P[8];
P[0]=wm;
for(int i=1;i<8;i++) P[i]=mont_s(P[i-1],wm,mod,ninv);
__m256i Pv=_mm256_loadu_si256((__m256i*)P);
__m256i carry=_mm256_set1_epi32((int)tmp[0]);
int p=1;
for(;p+8<=s;p+=8){
__m256i chunk=mont8(carry,Pv);
_mm256_storeu_si256((__m256i*)(tmp+p),chunk);
carry=_mm256_set1_epi32((int)tmp[p+7]);
}
for(;p<s;p++) tmp[p]=mont_s(tmp[p-1],wm,mod,ninv);
}
// Forward twiddles: TW[0..N/2) = w^p (radix-2); then 7 radix-4 stages interleaved [w1|w2|w3] per 8 groups.
static void gen_tw4(u32* tw, u32 w, u32 mod, u32 ninv, u32 r2c) {
gen_powers(tw, SIZE/2, mont_s(w,r2c,mod,ninv), mod, ninv, r2c);
int off=SIZE/2;
for(int len=SIZE/2; len>=32; len>>=2){
int s=len>>2; int step=SIZE/len;
u32 ws=powmod(w,step,mod); u32 wm=mont_s(ws,r2c,mod,ninv);
gen_powers(tmp,s,wm,mod,ninv,r2c);
for(int p=0;p<s;p+=8){
__m256i t1=loadu(tmp+p);
__m256i t2=mont8(t1,t1);
__m256i t3=mont8(t2,t1);
u32* wb=tw+off+p*3;
storeu(wb+0,t1); storeu(wb+8,t2); storeu(wb+16,t3);
}
off+=3*s;
}
}
// Inverse twiddles: 7 radix-4 stages (iw) then radix-2 stage iw^p.
static void gen_itw4(u32* tw, u32 iw, u32 mod, u32 ninv, u32 r2c) {
int off=0;
for(int len=32; len<=SIZE/2; len<<=2){
int s=len>>2; int step=SIZE/len;
u32 ws=powmod(iw,step,mod); u32 wm=mont_s(ws,r2c,mod,ninv);
gen_powers(tmp,s,wm,mod,ninv,r2c);
for(int p=0;p<s;p+=8){
__m256i t1=loadu(tmp+p);
__m256i t2=mont8(t1,t1);
__m256i t3=mont8(t2,t1);
u32* wb=tw+off+p*3;
storeu(wb+0,t1); storeu(wb+8,t2); storeu(wb+16,t3);
}
off+=3*s;
}
gen_powers(tw+off, SIZE/2, mont_s(iw,r2c,mod,ninv), mod, ninv, r2c);
}
// Vectorized len=8 DFT (base stage): 8 groups x 8 points per iter (reused from radix-8).
static void ntt8_vec(u32* x, __m256i z1, __m256i z3, __m256i iv) {
for(int i=0;i<SIZE;i+=64){
__m256i a0=loadu(x+i+0), a1=loadu(x+i+8), a2=loadu(x+i+16), a3=loadu(x+i+24);
__m256i a4=loadu(x+i+32), a5=loadu(x+i+40), a6=loadu(x+i+48), a7=loadu(x+i+56);
__m256i b0=_mm256_unpacklo_epi32(a0,a1), b1=_mm256_unpackhi_epi32(a0,a1);
__m256i b2=_mm256_unpacklo_epi32(a2,a3), b3=_mm256_unpackhi_epi32(a2,a3);
__m256i b4=_mm256_unpacklo_epi32(a4,a5), b5=_mm256_unpackhi_epi32(a4,a5);
__m256i b6=_mm256_unpacklo_epi32(a6,a7), b7=_mm256_unpackhi_epi32(a6,a7);
__m256i c0=_mm256_unpacklo_epi64(b0,b2), c1=_mm256_unpackhi_epi64(b0,b2);
__m256i c2=_mm256_unpacklo_epi64(b1,b3), c3=_mm256_unpackhi_epi64(b1,b3);
__m256i c4=_mm256_unpacklo_epi64(b4,b6), c5=_mm256_unpackhi_epi64(b4,b6);
__m256i c6=_mm256_unpacklo_epi64(b5,b7), c7=_mm256_unpackhi_epi64(b5,b7);
__m256i t0=_mm256_permute2f128_si256(c0,c4,0x20), t1=_mm256_permute2f128_si256(c1,c5,0x20);
__m256i t2=_mm256_permute2f128_si256(c2,c6,0x20), t3=_mm256_permute2f128_si256(c3,c7,0x20);
__m256i t4=_mm256_permute2f128_si256(c0,c4,0x31), t5=_mm256_permute2f128_si256(c1,c5,0x31);
__m256i t6=_mm256_permute2f128_si256(c2,c6,0x31), t7=_mm256_permute2f128_si256(c3,c7,0x31);
__m256i d0=addm(t0,t4), d1=addm(t1,t5), d2=addm(t2,t6), d3=addm(t3,t7);
__m256i e0=subm(t0,t4);
__m256i e1=mont8(subm(t1,t5),z1v);
__m256i e2=mont8(subm(t2,t6),ivv);
__m256i e3=mont8(subm(t3,t7),z3v);
__m256i s02=addm(d0,d2), t02=subm(d0,d2);
__m256i s13=addm(d1,d3), t13=mont8(subm(d1,d3),ivv);
__m256i D0=addm(s02,s13), D2=subm(s02,s13);
__m256i D1=addm(t02,t13), D3=subm(t02,t13);
__m256i u02=addm(e0,e2), v02=subm(e0,e2);
__m256i u13=addm(e1,e3), v13=mont8(subm(e1,e3),ivv);
__m256i E0=addm(u02,u13), E2=subm(u02,u13);
__m256i E1=addm(v02,v13), E3=subm(v02,v13);
__m256i rb0=_mm256_unpacklo_epi32(D0,E0), rb1=_mm256_unpackhi_epi32(D0,E0);
__m256i rb2=_mm256_unpacklo_epi32(D1,E1), rb3=_mm256_unpackhi_epi32(D1,E1);
__m256i rb4=_mm256_unpacklo_epi32(D2,E2), rb5=_mm256_unpackhi_epi32(D2,E2);
__m256i rb6=_mm256_unpacklo_epi32(D3,E3), rb7=_mm256_unpackhi_epi32(D3,E3);
__m256i rc0=_mm256_unpacklo_epi64(rb0,rb2), rc1=_mm256_unpackhi_epi64(rb0,rb2);
__m256i rc2=_mm256_unpacklo_epi64(rb1,rb3), rc3=_mm256_unpackhi_epi64(rb1,rb3);
__m256i rc4=_mm256_unpacklo_epi64(rb4,rb6), rc5=_mm256_unpackhi_epi64(rb4,rb6);
__m256i rc6=_mm256_unpacklo_epi64(rb5,rb7), rc7=_mm256_unpackhi_epi64(rb5,rb7);
__m256i r0=_mm256_permute2f128_si256(rc0,rc4,0x20), r1=_mm256_permute2f128_si256(rc1,rc5,0x20);
__m256i r2=_mm256_permute2f128_si256(rc2,rc6,0x20), r3=_mm256_permute2f128_si256(rc3,rc7,0x20);
__m256i r4=_mm256_permute2f128_si256(rc0,rc4,0x31), r5=_mm256_permute2f128_si256(rc1,rc5,0x31);
__m256i r6=_mm256_permute2f128_si256(rc2,rc6,0x31), r7=_mm256_permute2f128_si256(rc3,rc7,0x31);
storeu(x+i+0,r0); storeu(x+i+8,r1); storeu(x+i+16,r2); storeu(x+i+24,r3);
storeu(x+i+32,r4); storeu(x+i+40,r5); storeu(x+i+48,r6); storeu(x+i+56,r7);
}
}
// Forward radix-4 DIF: one 8-group chunk (4 loads, 4 mont8, 4 stores).
static inline void dif4_fwd8(u32* y, int s, const u32* wb) {
__m256i X0=loadu(y);
__m256i X1=loadu(y+s);
__m256i X2=loadu(y+2*s);
__m256i X3=loadu(y+3*s);
__m256i w1=loadu(wb);
__m256i w2=loadu(wb+8);
__m256i w3=loadu(wb+16);
__m256i t0=addm(X0,X2);
__m256i t1=subm(X0,X2);
__m256i t2=addm(X1,X3);
__m256i t3=mont8(subm(X1,X3),ivv);
__m256i D0=addm(t0,t2);
__m256i D2=subm(t0,t2);
__m256i D1=addm(t1,t3);
__m256i D3=subm(t1,t3);
storeu(y,D0);
storeu(y+s,mont8(D1,w1));
storeu(y+2*s,mont8(D2,w2));
storeu(y+3*s,mont8(D3,w3));
}
// Inverse radix-4 DIT: one 8-group chunk.
static inline void dit4_fwd8(u32* y, int s, const u32* wb) {
__m256i X0=loadu(y);
__m256i X1=loadu(y+s);
__m256i X2=loadu(y+2*s);
__m256i X3=loadu(y+3*s);
__m256i w1=loadu(wb);
__m256i w2=loadu(wb+8);
__m256i w3=loadu(wb+16);
__m256i x1=mont8(X1,w1);
__m256i x2=mont8(X2,w2);
__m256i x3=mont8(X3,w3);
__m256i t0=addm(X0,x2);
__m256i t1=subm(X0,x2);
__m256i t2=addm(x1,x3);
__m256i t3=mont8(subm(x1,x3),ivv);
__m256i o0=addm(t0,t2);
__m256i o2=subm(t0,t2);
__m256i o1=addm(t1,t3);
__m256i o3=subm(t1,t3);
storeu(y,o0);
storeu(y+s,o1);
storeu(y+2*s,o2);
storeu(y+3*s,o3);
}
static void ntt_fwd(u32* x) {
// radix-2 DIF
{
int s=SIZE>>1;
for(int p=0;p<s;p+=8){
__m256i av=loadu(x+p);
__m256i bv=loadu(x+p+s);
__m256i wv=loadu(TW+p);
storeu(x+p,addm(av,bv));
storeu(x+p+s,mont8(subm(av,bv),wv));
}
}
// 7 radix-4 DIF stages
int off=SIZE>>1;
for(int len=SIZE>>1; len>=32; len>>=2){
int s=len>>2;
for(int i=0;i<SIZE;i+=len){
u32* y=x+i;
const u32* wb=TW+off;
for(int p=0;p<s;p+=8) dif4_fwd8(y+p,s,wb+p*3);
}
off+=3*s;
}
ntt8_vec(x, z1v, z3v, ivv);
}
static void ntt_inv(u32* x) {
ntt8_vec(x, z1v, z3v, ivv);
int off=0;
for(int len=32; len<=SIZE/2; len<<=2){
int s=len>>2;
for(int i=0;i<SIZE;i+=len){
u32* y=x+i;
const u32* wb=ITW+off;
for(int p=0;p<s;p+=8) dit4_fwd8(y+p,s,wb+p*3);
}
off+=3*s;
}
{
int s=SIZE>>1;
const u32* wb=ITW+off;
for(int p=0;p<s;p+=8){
__m256i av=loadu(x+p);
__m256i bv=mont8(loadu(x+p+s),loadu(wb+p));
storeu(x+p,addm(av,bv));
storeu(x+p+s,subm(av,bv));
}
}
}
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-9>=sa){const char* q=pos-9;
u32 v=(u32)(q[0]-'0'); v=v*10+(u32)(q[1]-'0'); v=v*10+(u32)(q[2]-'0'); v=v*10+(u32)(q[3]-'0');
v=v*10+(u32)(q[4]-'0'); v=v*10+(u32)(q[5]-'0'); v=v*10+(u32)(q[6]-'0'); v=v*10+(u32)(q[7]-'0'); v=v*10+(u32)(q[8]-'0');
limbsA[na++]=v; pos=q;
}
if(pos>sa){u32 v=0; for(const char* q=sa;q<pos;q++) v=v*10+(u32)(*q-'0'); limbsA[na++]=v;}
}
{
const char* pos=b_end;
while(pos-9>=sb){const char* q=pos-9;
u32 v=(u32)(q[0]-'0'); v=v*10+(u32)(q[1]-'0'); v=v*10+(u32)(q[2]-'0'); v=v*10+(u32)(q[3]-'0');
v=v*10+(u32)(q[4]-'0'); v=v*10+(u32)(q[5]-'0'); v=v*10+(u32)(q[6]-'0'); v=v*10+(u32)(q[7]-'0'); v=v*10+(u32)(q[8]-'0');
limbsB[nb++]=v; pos=q;
}
if(pos>sb){u32 v=0; for(const char* q=sb;q<pos;q++) v=v*10+(u32)(*q-'0'); limbsB[nb++]=v;}
}
shufvec=_mm256_setr_epi8(0,1,2,3,8,9,10,11,0,0,0,0,0,0,0,0, 0,1,2,3,8,9,10,11,0,0,0,0,0,0,0,0);
for(int mi=0;mi<3;mi++){
u32 mod=MODS[mi],ninv=NINVS[mi],r2c=R2S[mi];
curmod=mod; curninv=ninv;
modvec=_mm256_set1_epi32((int)mod); ninvvec=_mm256_set1_epi32((int)ninv);
u32 w=powmod(3u,(mod-1)/SIZE,mod);
u32 iw=powmod(w,mod-2,mod);
u32 Iv=mont_s(IROOT[mi],r2c,mod,ninv);
u32 Ivi=mont_s(IMINUS[mi],r2c,mod,ninv);
u32 Z1=mont_s(ZETA[mi],r2c,mod,ninv), Z3=mont_s(ZETA3[mi],r2c,mod,ninv);
u32 Z1i=mont_s(ZETA7[mi],r2c,mod,ninv), Z3i=mont_s(ZETA5[mi],r2c,mod,ninv);
gen_tw4(TW,w,mod,ninv,r2c);
gen_itw4(ITW,iw,mod,ninv,r2c);
__m256i r2b=_mm256_set1_epi32((int)r2c);
{int vlim=(na+7)>>3;for(int v=0;v<vlim;v++){__m256i x=loadu(limbsA+v*8);x=_mm256_min_epu32(x,_mm256_sub_epi32(x,modvec));x=_mm256_min_epu32(x,_mm256_sub_epi32(x,modvec));storeu(A+v*8,mont8(x,r2b));}memset(A+vlim*8,0,(SIZE-vlim*8)*4);}
{int vlim=(nb+7)>>3;for(int v=0;v<vlim;v++){__m256i x=loadu(limbsB+v*8);x=_mm256_min_epu32(x,_mm256_sub_epi32(x,modvec));x=_mm256_min_epu32(x,_mm256_sub_epi32(x,modvec));storeu(B+v*8,mont8(x,r2b));}memset(B+vlim*8,0,(SIZE-vlim*8)*4);}
z1v=_mm256_set1_epi32((int)Z1); z3v=_mm256_set1_epi32((int)Z3); ivv=_mm256_set1_epi32((int)Iv);
ntt_fwd(A);
ntt_fwd(B);
for(int v=0;v<SIZE;v+=8) storeu(A+v,mont8(loadu(A+v),loadu(B+v)));
z1v=_mm256_set1_epi32((int)Z1i); z3v=_mm256_set1_epi32((int)Z3i); ivv=_mm256_set1_epi32((int)Ivi);
ntt_inv(A);
u32 ninvn=powmod(SIZE,mod-2,mod);
__m256i ninvb=_mm256_set1_epi32((int)ninvn);
u32* dst=(mi==0)?r0:(mi==1)?r1:r2;
for(int v=0;v<SIZE;v+=8) storeu(dst+v,mont8(loadu(A+v),ninvb));
}
{
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_128=(u128)M0*M1; u64 M01=(u64)M01_128;
u64 INV012=powmod(M01%m2,m2-2,m2);
u64 M0m2=M0%m2;
u64 M01_q=M01/1000000000u; u64 M01_r=M01%1000000000u;
int outlen=na+nb-1; u64 carry=0; u64 m2x2=2ull*m2;
for(int i=0;i<outlen;i++){
u64 a0=r0[i],a1=r1[i],a2=r2[i];
u64 t1=a1-a0+m1; if(t1>=m1) t1-=m1;
t1=t1*INV01%m1;
u64 a0m2=a0; if(a0m2>=m2x2) a0m2-=m2x2; if(a0m2>=m2) a0m2-=m2;
u64 x1m2=(a0m2+t1*M0m2)%m2;
u64 t2=a2-x1m2+m2; if(t2>=m2) t2-=m2;
t2=t2*INV012%m2;
u64 d=a0+t1*M0+carry;
u64 s=d+t2*M01_r;
u64 carr=t2*M01_q+s/1000000000u;
outlimbs[i]=(u32)(s%1000000000u); carry=carr;
}
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
}
| Compilation | N/A | N/A | Compile OK | Score: N/A | 显示更多 |
| Testcase #1 | 15.884 ms | 10 MB + 764 KB | Accepted | Score: 100 | 显示更多 |