洛谷 P8367 [LNOI2022] 盒 题解
题目链接：https://www.luogu.com.cn/problem/P8367
#include <bits/stdc++.h>
using namespace std;
const int p=998244353;
long long T,n,S,a[500005],w[500005],fac[3000005],ifac[3000005],c[500005];
inline long long qpow(long long a,long long b){
int res=1;
for (;b;b>>=1,a=1ll*a*a%p) if (b&1) res=1ll*res*a%p;
return res;
}
inline long long C(long long n,long long m){
return fac[n]*ifac[m]%p*ifac[n-m]%p;
}
struct f{
long long n,S,i,k,res;
void init(long long n0,long long S0,long long i0,long long k0){
n=n0,S=S0,i=i0,k=k0,res=0;
for (long long di=0;di<=k;di++) (res+=1ll*C(di+i-1,i-1)*C(S-di+n-i-1,n-i-1)%p)%=p;
}
void movek(long long K){
if (k>=K) return;
for (;k<K;k++) res+=C(k+i,i-1)*C(S-k+n-i-2,n-i-1)%p;
res%=p;
}
void movei(long long I){
if (i>=I) return;
for (;i<I;i++) res+=p-C(i+k,i)*C(n+S-i-k-2,S-k-1)%p;
res%=p;
}
}f1,f2,f3,f4;
int main(){
scanf("%lld",&T);
fac[0]=1;
for (int i=1;i<=3e6;i++) fac[i]=fac[i-1]*i%p;
ifac[3000000]=qpow(fac[3000000],p-2);
for (int i=3e6-1;i>=0;i--) ifac[i]=ifac[i+1]*(i+1)%p;
while (T--){
scanf("%lld",&n),S=0;
for (int i=1;i<=n;i++) scanf("%lld",&a[i]),S+=a[i],c[i]=c[i-1]+a[i];
for (int i=1;i<n;i++) scanf("%lld",&w[i]);
long long ans=0;
f1.init(n,S,1,c[1]),f2.init(n,S,1,S),f3.init(n+1,S-1,2,S-1),f4.init(n+1,S-1,2,c[1]-1);
for (int i=1;i<n;i++){
f1.movei(i),f2.movei(i),f3.movei(i+1),f4.movei(i+1);
f1.movek(c[i]),f4.movek(c[i]-1);
(ans+=w[i]*(c[i]*(2*f1.res-f2.res+p)%p+i*(f3.res-2*f4.res+p)%p))%=p;
}
printf("%lld\n",(ans+p)%p);
}
return 0;
}