牛客练习赛50 F.tokitsukaze and Another Protoss and Zerg(启发式NTT)(模板题)

题目

在n场1v1中，假设第i场，有ai个星灵单位，和bi个虫族单位。

tokitsukaze可以在一场1v1中，任选一种种族进行游戏。

如果选择了星灵，那么在这场游戏中，可以选择出兵1到ai个单位。

那么同理，如果选择了虫族，那么在这场游戏中，可以选择出兵1到bi个单位。
tokitsukaze想知道，恰好选择了0,1,2,3,...,个星灵单位的方案数分别是多少。

由于答案很大，所以输出答案mod 998244353 后的结果。

假设所有星灵单位互不相同，所有异虫单位也互不相同。

注意：若两个方案，有其中一个单位不同，即视为不相同。

n<=1e5，ai<=2e5，bi<=1e5，。

题解

考虑生成函数为，

c0即选b[i]的非空子集，方案数

其余ck(1<=k<=ai)即从a[i]中选k个，方案数

放入系数之后，直接启发式NTT即可，

开始试了一发分治+FFT炸了精度

代码

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N=1<<18,mod=998244353,G=3;
 
int n,sum;
int a[N],b[N],fac[N],finv[N];

 
ll A[N],B[N],inv;
int R[N];
int power(int x,int n)
{
	int res=1;
	for(;n;n>>=1,x=1ll*x*x%mod)
	if(n&1)res=1ll*res*x%mod;
	return res;
}
vector <int> colors[N<<1];
struct cmp{
	bool operator ()(int a,int b){
		return colors[a].size()>colors[b].size();
	}
};
priority_queue <int,vector<int>,cmp> heap;
void NTT(ll a[],int n,int re){
	for (int i=0;i<n;i++)
		if (i<R[i])
			swap(a[i],a[R[i]]);
    for(int i=2;i<=n;i<<=1)
    {
        ll mid=i>>1;
        ll wn=power(G,(mod-1)/i);
        if(re) wn=power(wn,(mod-2));
        for(int j=0;j<n;j+=i)
        {
            ll w=1;
            for(int k=0;k<mid;k++)
            {
                ll temp1=a[j+k];
                ll temp2=a[j+k+mid]*w%mod;
                a[j+k]=(temp1+temp2)%mod;
                a[j+k+mid]=(temp1-temp2);
				a[j+k+mid]=(a[j+k+mid]%mod+mod)%mod;
                w=w*wn%mod;
            }
        }
    }
    if(re){
		for(int i=0;i<n;i++)
       	 	a[i]=(a[i]*inv)%mod;
    }
}
void NTT_times(vector <int> &a,vector <int> &b,vector <int> &c){
	int n,d;
	for (int i=0;i<a.size();i++)
		A[i]=a[i];
	for (int i=0;i<b.size();i++)
		B[i]=b[i];
	for (n=1,d=0;n<a.size()+b.size()-1;n<<=1,d++);
	inv=power(n,mod-2);
	for (int i=0;i<n;i++){
		R[i]=(R[i>>1]>>1)|((i&1)<<(d-1));
	}
	for (int i=a.size();i<n;i++)
		A[i]=0;
	for (int i=b.size();i<n;i++)
		B[i]=0;
	NTT(A,n,0),NTT(B,n,0);
	for (int i=0;i<n;i++)
		A[i]=A[i]*B[i]%mod;
	NTT(A,n,1);
	c.clear();
	for (int i=0;i<=a.size()+b.size()-2;i++)
		c.push_back(A[i]%mod);
}
void init()
{
	fac[0]=finv[0]=1;
	for(int i=1;i<N;++i)
	fac[i]=1ll*fac[i-1]*i%mod;
	finv[N-1]=power(fac[N-1],mod-2);
	for(int i=N-2;i>=1;--i)
	finv[i]=1ll*finv[i+1]*(i+1)%mod;
}
int C(int n,int m)
{
	if(m<0||n<m)return 0;
	return 1ll*fac[n]*finv[m]%mod*finv[n-m]%mod;
}
int main(){
	init();
	scanf("%d",&n);
	for(int i=1;i<=n;++i)
	{
		scanf("%d",&a[i]);
		sum+=a[i];
	}
	for(int i=1;i<=n;++i)
	scanf("%d",&b[i]);
	while (!heap.empty())
		heap.pop();
	int sz=0;
	for (int i=1;i<=n;i++){
		colors[++sz].clear();
		colors[sz].push_back(power(2,b[i])-1);//选b[i]的一个非空子集 
		for(int j=1;j<=a[i];++j)
		colors[sz].push_back(C(a[i],j));//a[i]里面选j个
		heap.push(sz);
	}
	while (heap.size()>=2){
		int x=heap.top();
		heap.pop();
		int y=heap.top();
		heap.pop();
		NTT_times(colors[x],colors[y],colors[++sz]);
		colors[x].clear(),colors[y].clear();
		heap.push(sz);
	}
	for(int i=0;i<=sum;++i)
	printf("%d%c",colors[sz][i],i==sum?'\n':' ');
	return 0;
}