2022牛客多校#4 C. Easy Counting Problem

C. Easy Counting Problem

题目大意

定义好字符串，满足：

• 包含十进制下小于 $w(2\le w\le 10)$ 的数字
• 数字 $i$ 至少出现 $c_i(1\le c_i\le 50000,\sum c_i\le 50000)$ 次

有 $q(1\le q\le 300)$ 次询问，每次询问求解有多少个不同的长度为 $n(1\le n\le 10^7)$ 的好字符串。

题解

考虑生成函数，由于求排列计数，因此考虑EGF。
数字 $k$ 至少出现 $c_k$ 次，其对应的EGF为
$$\begin{array}{rl}{f}_k(x)& =\frac{x^{c_k}}{c_k!}+\frac{x^{c_k+1}}{(c_k+1)!}+\frac{x^{c_k+2}}{(c_k+2)!}+...\\ & =e^x-\sum _{i=0}^{c_k-1}\frac{x^i}{i!}\end{array}$$

所有条件汇总的EGF为
$$F(x)=\prod _{k=0}^{w-1}{f}_k(x)=\prod _{k=0}^{w-1}(e^x-\sum _{i=0}^{c_k-1}\frac{x^i}{i!})$$

对于每个询问的 $n$ ，$F(x)$ 的 $[\frac{x^n}{n!}]$ 的系数即为答案。
$$an{s}_n=[\frac{x^n}{n!}]\prod _{k=0}^{w-1}(e^x-\sum _{i=0}^{c_k-1}\frac{x^i}{i!})$$

上式无法直接卷积，考虑单独提出 $e^x$
$$F(x)=\sum _{i=0}^{w-1}{e}^{ix}\cdot (-1{)}^{w-1-i}\cdot g_{w-1,w-1-i}(x)$$

其中 $g_{i,j}$ 表示在 { ∑ k = 0 c i − 1 x k k ! } \{\sum_{k=0}^{c_i-1}{\frac{x^k}{k!}}\} {∑k=0ci​−1​k!xk​} 的前 $i$ 项中，选取了 $j$ 项之积的表达式之和。存在转移式
$$g_{i,j}=g_{i-1,j}+g_{i-1,j-1}\cdot \sum _{k=0}^{c_i-1}\frac{x^k}{k!}$$

可以采用NTT在 $O(w^{2}C\log C)$ 的复杂度内（其中 $C=\sum c_i$），求解出所有的 $g_{i,j}(x)$ ，并且其最高次不超过 $C$ 。

考虑拆开 $e^{ix}={\sum }_{j=0}^{+\infty }\frac{i^j}{j!}{x}^j$ ，
$$\begin{array}{rl}F(x)& =\sum _{i=0}^{w-1}(-1{)}^{w-1-i}\cdot g_{w-1,w-1-i}(x)\sum _{j=0}^{+\infty }\frac{i^j}{j!}{x}^j\\ & =\sum _{i=0}^{w-1}{h}_i(x)\sum _{j=0}^{+\infty }\frac{i^j}{j!}{x}^j\end{array}$$

求解上式中 $[\frac{x^n}{n!}]$ 的系数，由于最高次 $C=\sum c_i\le 50000$ ，依次可以枚举 $h_i(x)$ 中第 $k$ 项 $x^k$ 的系数 $[x^k]h_i(x)$ ，则还需在 ${\sum }_{j=0}^{+\infty }\frac{i^j}{j!}{x}^j$ 中选取 $n-k$ 项的系数，即 $\frac{i^{n-j}}{(n-j)!}$

因此答案为
$$an{s}_n=n!\cdot \sum _{i=0}^{w-1}\sum _{k=0}^n\frac{i^{n-j}}{(n-j)!}[x^k]h_i(x)$$

其中 $n!$ 用于将 $x^n$ 的系数转化为 $[\frac{x^n}{n!}]$ 的系数。

参考代码

来自SolitaryDrea

#include<bits/stdc++.h>
using namespace std;
#define FOR(i,s,t) for(int i=(s),_t=(t); i<=_t; ++i)
#define DOR(i,s,t) for(int i=(s),_t=(t); i>=_t; --i)
typedef long long ll;
const int Mod=998244353;
const int N=2e5+50;
const int M=1e7+50;
bool __;

ll Fast(ll x,int b) {
	ll t=1;
	for(; b; b>>=1,x=x*x%Mod) {
		if(b&1) t=t*x%Mod;
	}
	return t;
}
void DFT(ll *a,int n,int f) {
	static int rev[N],ww[N],iw[N],pn=0;
	if(pn!=n) {
		ll p=Fast(3,(Mod-1)/n);
		ww[0]=1;
		FOR(i,1,n-1) ww[i]=ww[i-1]*p%Mod;
		iw[n-1]=Fast(ww[n-1],Mod-2);
		DOR(i,n-1,1) iw[i-1]=iw[i]*p%Mod;
		FOR(i,0,n-1) rev[i]=(rev[i>>1]>>1)|((i&1)*(n>>1));
		pn=n;
	}
	int *w=(f>0)?ww:iw;
	FOR(i,0,n-1) if(rev[i]<i) swap(a[rev[i]],a[i]);
	for(int l=2; l<=n; l<<=1) {
		for(ll *p=a; p!=a+n; p+=l) {
			for(int i=0,m=l>>1; i<m; ++i) {
				ll t=p[i+m];
				p[i+m]=(p[i]+w[n/l*(i+m)]*t)%Mod;
				p[i]=(p[i]+w[n/l*i]*t)%Mod;
			}
		}
	}
	if(f<0) {
		ll t=Fast(n,Mod-2);
		FOR(i,0,n-1) a[i]=a[i]*t%Mod;
	}
}
void Poly_Mul(const ll *A,int LenA,const ll *B,int LenB,ll *C) {
	static ll a[N],b[N];
	int n=1;for(; n<LenA+LenB; n<<=1);
	FOR(i,0,n-1) a[i]=b[i]=0;
	FOR(i,0,LenA-1) a[i]=A[i];
	FOR(i,0,LenB-1) b[i]=B[i];
	DFT(a,n,1);DFT(b,n,1);
	FOR(i,0,n-1) C[i]=a[i]*b[i]%Mod;
	DFT(C,n,-1);
}

ll Fac[M],Inv[M];
ll a[12][N];
ll b[N],g[N];
ll C(int n,int m) {
	return Fac[n]*Inv[n-m]%Mod*Inv[m]%Mod;
}
bool ___;
int main() {
	int n=1e7;
	Fac[0]=1;
	FOR(i,1,n) Fac[i]=Fac[i-1]*i%Mod;
	Inv[n]=Fast(Fac[n],Mod-2);
	DOR(i,n,1) Inv[i-1]=Inv[i]*i%Mod;
	int w;
	scanf("%d",&w);
	a[0][0]=1;
	int K=0;
	FOR(i,1,w) {
		int c;
		scanf("%d",&c);
		FOR(i,0,c-1) b[i]=-Inv[i];
		DOR(j,i-1,0) {
			Poly_Mul(a[j],K+1,b,c,g);
			FOR(k,0,K+c-1) a[j+1][k]=(a[j+1][k]+g[k])%Mod;
		}
		
		K+=c-1;
	}
	int q;
	scanf("%d",&q);
	while(q--) {
		scanf("%d",&n);
		ll res=0;
		FOR(i,0,w) {
			ll u=Fast(w-i,n);
			ll v=Fast(w-i,Mod-2);
			FOR(j,0,min(K,n)) {
				if(w-i) {
					res=(res+Fac[n]*a[i][j]%Mod*Inv[n-j]%Mod*u)%Mod;
					u=u*v%Mod;
				} else if(n-j==0) {
					res=(res+Fac[n]*a[i][j]%Mod*Inv[n-j]%Mod)%Mod;
				}
			}
		}
		if(res<0) res+=Mod;
		printf("%lld\n",res);
	}
	return 0;
}