CodeChef 2020 November - Challenge Chef and the Combination Lock (多项式)

CodeChef 2020 November - Challenge Chef and the Combination Lock (多项式)

题目大意：给定了\(n\)个随机变量\(x_i\in{0,1,\cdots,A_i}\)，令\(\Chi=\min_i\lbrace x_i,A_i-x_i\rbrace\)，求\(E(\Chi)\)

我们知道\(E(\Chi)=\sum_{i=0}^{\infty} P(\Chi>i)\)

不妨考虑计算\(P(\Chi>i)\)，先计算方案数，发现方案数可以用一个多项式来表示

令\(F(x)\)为\(\Chi>x\)的方案数，则\(\begin{aligned}F(x)=\prod_{i=1}^n (A_i-1-2x)　\end{aligned}\)

显然\(\Chi \leq \min\lbrace \frac{A_i}{2} \rbrace\)，不妨设这个上界为\(U\)，也就是说我们要求\(\sum_{i=0}^U F(i)\)

常识：一个\(n\)次多项式前缀和可以用一个不超过\(n+1\)次的多项式来表示

如果暴力求出\(F(x)\)在\(x=0,1,\cdots,n+1\)处的值，累前缀和，然后用拉格朗日插值法求出解

暴力求值复杂度为\(O(n^2)\)，拉格朗日插值复杂度为\(O(n)\)

可以用分治\(\text{NTT}\)优化\(F(x)\)的求解，然后用多项式多点求值求得点值

复杂度为\(O(n\log ^2n)\)，实际在CodeChef上的运行时间为0.53s

#include<bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef double db;
typedef unsigned long long ull;
typedef pair <int,int> Pii;
#define reg register
#define pb push_back
#define mp make_pair
#define Mod1(x) ((x>=P)&&(x-=P))
#define Mod2(x) ((x<0)&&(x+=P))
#define rep(i,a,b) for(int i=a,i##end=b;i<=i##end;++i)
#define drep(i,a,b) for(int i=a,i##end=b;i>=i##end;--i)
template <class T> inline void cmin(T &a,T b){ ((a>b)&&(a=b)); }
template <class T> inline void cmax(T &a,T b){ ((a<b)&&(a=b)); }

char IO;
template <class T=int> T rd(){
	T s=0; int f=0;
	while(!isdigit(IO=getchar())) if(IO=='-') f=1;
	do s=(s<<1)+(s<<3)+(IO^'0');
	while(isdigit(IO=getchar()));
	return f?-s:s;
}

const int N=1<<18,P=998244353;


ll qpow(ll x,ll k=P-2) {
	ll res=1;
	for(;k;k>>=1,x=x*x%P) if(k&1) res=res*x%P;
	return res;
}

typedef vector <int> Poly;
void Show(Poly a,int k=0){ 
	if(!k){ for(int i:a) printf("%d ",i); puts(""); }
	else for(int i:a) printf("%d\n",i);
}
int rev[N],w[N];
int Inv[N+1],Fac[N+1],FInv[N+1];

void Init() { 
	int t=qpow(3,(P-1)/N);
	w[N>>1]=1;
	rep(i,(N>>1)+1,N-1) w[i]=1ll*w[i-1]*t%P;
	drep(i,(N>>1)-1,1) w[i]=w[i<<1];
	Inv[0]=Inv[1]=Fac[0]=Fac[1]=FInv[0]=FInv[1]=1;
	rep(i,2,N) {
		Inv[i]=1ll*(P-P/i)*Inv[P%i]%P; 
		FInv[i]=1ll*FInv[i-1]*Inv[i]%P;
		Fac[i]=1ll*Fac[i-1]*i%P;
	}

}
int Init(int n){
	int R=1,c=-1;
	while(R<n) R<<=1,c++;
	rep(i,1,R-1) rev[i]=(rev[i>>1]>>1)|((i&1)<<c);
	return R;
}


void NTT(int n,int *a,int f){
	rep(i,0,n-1) if(rev[i]<i) swap(a[i],a[rev[i]]);
	for(int i=1;i<n;i<<=1) {
		int *e=w+i;
		for(int l=0;l<n;l+=i*2) {
			for(int j=l;j<l+i;++j) {
				int t=1ll*a[j+i]*e[j-l]%P;
				a[j+i]=a[j]-t,Mod2(a[j+i]);
				a[j]+=t,Mod1(a[j]);
			}
		}
	}
	if(f==-1) {
		reverse(a+1,a+n);
		ll base=Inv[n];
		rep(i,0,n-1) a[i]=a[i]*base%P;
	}
}
void NTT(int n,Poly &a,int f){
	static int A[N];
	if((int)a.size()<n) a.resize(n);
	rep(i,0,n-1) A[i]=a[i];
	NTT(n,A,f);
	rep(i,0,n-1) a[i]=A[i];
}

Poly operator * (Poly a,Poly b){
	int n=a.size()+b.size()-1;
	int R=Init(n);
	a.resize(R),b.resize(R);
	NTT(R,a,1),NTT(R,b,1);
	rep(i,0,R-1) a[i]=1ll*a[i]*b[i]%P;
	NTT(R,a,-1);
	a.resize(n);
	return a;
}

Poly operator + (Poly a,Poly b) { 
	int n=max(a.size(),b.size());
	a.resize(n),b.resize(n);
	rep(i,0,n-1) a[i]+=b[i],Mod1(a[i]);
	return a; 
}
Poly operator - (Poly a,Poly b) { 
	int n=max(a.size(),b.size());
	a.resize(n),b.resize(n);
	rep(i,0,n-1) a[i]-=b[i],Mod2(a[i]);
	return a; 
}

Poly Poly_Inv(Poly a) { 
	int n=a.size();
	if(n==1) return Poly{(int)qpow(a[0],P-2)};
	Poly b=a; b.resize((n+1)/2); b=Poly_Inv(b);
	int R=Init(n<<1);
	a.resize(R),b.resize(R);
	NTT(R,a,1),NTT(R,b,1);
	rep(i,0,R-1) a[i]=(2-1ll*a[i]*b[i]%P+P)*b[i]%P;
	NTT(R,a,-1);
	a.resize(n);
	return a;
}

// 应用转置原理优化的多项式多点求值
Poly Evaluate(Poly F,Poly X){
	static int ls[N<<1],rs[N<<1],cnt;
	static Poly T[N<<1];
	static auto TMul=[&] (Poly F,Poly G){
		int n=F.size(),m=G.size();
		if(n<=20 && m<=20){
			rep(i,0,n-m) {
				int t=0;
				rep(j,0,m-1) t=(t+1ll*F[i+j]*G[j])%P;
				F[i]=t;
			} 
			F.resize(n-m+1);
			return F;
		}
		reverse(G.begin(),G.end());
		int R=Init(n);
		NTT(R,F,1),NTT(R,G,1);
		rep(i,0,R-1) F[i]=1ll*F[i]*G[i]%P;
		NTT(R,F,-1); Poly T(n-m+1);
		rep(i,0,n-m) T[i]=F[i+m-1];
		return T;
	};
	static function <int(int,int)> Build=[&](int l,int r) {
		int u=++cnt; ls[u]=rs[u]=0;
		if(l==r) {
			T[u]=Poly{1,P-X[l]};
			return u;
		}
		int mid=(l+r)>>1;
		ls[u]=Build(l,mid),rs[u]=Build(mid+1,r);
		T[u]=T[ls[u]]*T[rs[u]];
		return u;
	};

	int n=F.size(),m=X.size();
	cmax(n,m),F.resize(n),X.resize(n);
	cnt=0,Build(0,n-1);
	F.resize(n*2+1),T[1]=TMul(F,Poly_Inv(T[1]));
	int p=0;
	rep(i,1,cnt) if(ls[i]) {
		swap(T[ls[i]],T[rs[i]]);

		int R=Init(T[i].size()),n=T[i].size(),m1=T[ls[i]].size(),m2=T[rs[i]].size();
		NTT(R,T[i],1);
		reverse(T[ls[i]].begin(),T[ls[i]].end()); reverse(T[rs[i]].begin(),T[rs[i]].end());
		NTT(R,T[ls[i]],1); NTT(R,T[rs[i]],1);
		rep(j,0,R-1) {
			T[ls[i]][j]=1ll*T[ls[i]][j]*T[i][j]%P;
			T[rs[i]][j]=1ll*T[rs[i]][j]*T[i][j]%P;
		}
		NTT(R,T[ls[i]],-1); NTT(R,T[rs[i]],-1);
		rep(j,0,n-m1) T[ls[i]][j]=T[ls[i]][j+m1-1];
		T[ls[i]].resize(n-m1+1);
		rep(j,0,n-m2) T[rs[i]][j]=T[rs[i]][j+m2-1];
		T[rs[i]].resize(n-m2+1);
	} else X[p++]=T[i][0];
	X.resize(m);
	return X;
}

int n;
int A[N];
int I[N],J[N];
int F[N],L[N],R[N];

Poly Solve(int l,int r) {
	if(l==r) return Poly{A[l]-1,P-2};
	int mid=(l+r)>>1;
	return Solve(l,mid)*Solve(mid+1,r);
}


int main() {
	Init();
	rep(i,J[0]=1,N-1) J[i]=1ll*J[i-1]*i%P;
	I[N-1]=qpow(J[N-1]);
	drep(i,N-1,1) I[i-1]=1ll*I[i]*i%P;

	rep(kase,1,rd()) {
		int x=P,All=1; 
		n=rd();
		rep(i,1,n+1) F[i]=0;
		F[0]=1;
		int f=0;
		rep(i,1,n) {
			A[i]=rd();
			f|=!A[i];
			cmin(x,(A[i]-1)/2),All=1ll*All*(A[i]+1)%P;
		}
		if(f){ puts("0"); continue; }
		Poly Y=Solve(1,n),X(n+2);
		rep(i,0,n+1) X[i]=i;
		Y=Evaluate(Y,X);
		rep(i,1,n+1) Y[i]=(Y[i]+Y[i-1])%P;
		int ans=0;
		if(x<=n+1) ans=Y[x];
		else {
            // 拉格朗日插值
			L[0]=x;
			rep(i,1,n+1) L[i]=1ll*L[i-1]*(x-i)%P;
			R[n+2]=1;
			drep(i,n+1,0) R[i]=1ll*R[i+1]*(x-i)%P;
			rep(i,0,n+1) {
				int t=1ll*Y[i]*(i?L[i-1]:1)%P*R[i+1]%P*I[i]%P*I[n+1-i]%P;
				if((n+1-i)&1) t=P-t;
				ans=(ans+t)%P;
			}
		}
		ans=ans*qpow(All)%P;
		printf("%d\n",ans);
	}
}