Loj #2320 -「清华集训 2017」生成树计数

\[\sum _{\sum v_i = n-2} \prod (a_i ^{v_i+1} * (v_i+1) ^ m /v_i!) *(\sum (v_i+1)^m) \]

将 \(\sum (v_i+1)^m\) 中的贡献分开算。

我们有两个生成函数。

第一个：

\[\sum _i a^{i+1} * (v_i+1)^m x^i / i! \]

第二个：

\[\sum _i a^{i+1} * (i+1)^{2m} x^i / i! \]

先将 \(\prod a\) 算到前面。

令：

\[A(x) =\sum x^i * (i+1)^{m} /i! \\ B(x) =\sum x^i * (i+1)^{2m} /i! \]

我们要算的是：

\[\sum _i B(a_ix)/A(a_ix) * \prod A(a_jx) \]

后面可以ln，exp。

然后发现泥算个等幂和就行了。

补充一下等幂和怎么算：
生成函数是：

\[\sum \frac 1 {1-a_ix} \]

因为

\[\ln' (f(x)) = \frac {f'(x)}{f(x)} \\ \ln' (1-a_ix) = \frac {-a_i} {1-a_ix} = -a_i - a_i^2 x \]

所以可以一个分治FFT算。

最终复杂度$n\log ^2n $

#include<bits/stdc++.h>
using namespace std;
const int N=5e5+5;
typedef long long ll;

const int mod=998244353;
int add(int a,int b){a+=b;return a>=mod?a-mod:a;}
int sub(int a,int b){a-=b;return a<0?a+mod:a;}
int mul(int a,int b){return (ll)a*b%mod;}
int qpow(int a,int b){int ret=1;for(;b;b>>=1,a=mul(a,a))if(b&1)ret=mul(ret,a);return ret;}
/*math*/

namespace Template_Poly{
	typedef vector<int> Poly;
	int rev[N];
	Poly Poly_add(Poly A,Poly B){
		A.resize(max(A.size(),B.size()));
		for(size_t i=0;i<B.size();i++)A[i]=add(A[i],B[i]);
		return A;
	}
	Poly Poly_sub(Poly A,Poly B){
		A.resize(max(A.size(),B.size()));
		for(size_t i=0;i<B.size();i++)A[i]=sub(A[i],B[i]);
		return A;
	}
	void DFT(int *t,int n,int type){
		int l=0;while(1<<l<n)++l;
		for(int i=0;i<n;i++)rev[i]=(rev[i>>1]>>1)|((i&1)<<(l-1));
		for(int i=0;i<n;i++)if(rev[i]>i)swap(t[rev[i]],t[i]);
		for(int step=1;step<n;step<<=1){
			int wn=qpow(3,(mod-1)/(step<<1));
			for(int i=0;i<n;i+=step<<1){
				int w=1;
				for(int k=0;k<step;k++,w=mul(w,wn)){
					int x=t[i+k],y=mul(t[i+k+step],w);
					t[i+k]=add(x,y),t[i+k+step]=sub(x,y);
				}
			}
		}
		if(type==1)return;
		for(int i=1;i<n-i;i++)swap(t[i],t[n-i]);
		int inv=qpow(n,mod-2);
		for(int i=0;i<n;i++)t[i]=mul(t[i],inv);
	}
	Poly NTT(Poly A,int n,Poly B,int m){
		static Poly res,PolA,PolB;
		PolA=A,PolB=B;
		int len=1;while(len < n+m)len<<=1;
		res.resize(len);
		PolA.resize(len),PolB.resize(len);
		DFT(&PolA[0],len,1);DFT(&PolB[0],len,1);
		for(int i=0;i<len;i++) res[i]= mul(PolA[i],PolB[i]);
		DFT(&res[0],len,-1);
		res.resize(n+m-1);
		return res;
	}
	Poly NTT(Poly A,Poly B){
		return NTT(A,A.size(),B,B.size());
	}
	Poly Poly_inv(Poly A,int n){
		if(n==1)return Poly(1,qpow(A[0],mod-2));
		int len=1<<((int)ceil(log2(n))+1);
		Poly x=Poly_inv(A,(n+1)>>1),y;
		x.resize(len),y.resize(len);
		for(int i=0;i<n;i++)y[i]=A[i];
		DFT(&x[0],len,1),DFT(&y[0],len,1);
		for(int i=0;i<len;i++)x[i]=mul(x[i],sub(2,mul(x[i],y[i])));
		DFT(&x[0],len,-1);
		x.resize(n);
		return x;
	}
	Poly Poly_inv(Poly A){
		return Poly_inv(A,A.size());
	}
	Poly Deri(Poly A){
		int n=A.size();
		for(int i=1;i<n;i++)A[i-1]=mul(A[i],i);
		A.resize(n-1);
		return A;
	}

	Poly Inte(Poly A){
		int n=A.size();
		A.resize(n+1);
		for(int i=n;i;i--)A[i]=mul(A[i-1],qpow(i,mod-2));
		A[0]=0;
		return A;
	}

	Poly ln(Poly A){
		int len=A.size();
		A=Inte(NTT(Deri(A),Poly_inv(A)));
		A.resize(len);
		return A;
	}

	Poly exp(Poly A,int n){
		if(n==1)return Poly(1,1);
		Poly x=exp(A,(n+1)>>1),y;
		x.resize(n);
		y=ln(x);
		for(int i=0;i<n;i++)y[i]=sub(A[i],y[i]);
		y[0]++;
		x=NTT(x,y);
		x.resize(n);
		return x;
	}
	Poly exp(Poly A){
		return exp(A,A.size());
	}

	Poly sqrt(Poly A,int n){
		if(n==1)return Poly(1,1);
		Poly x=sqrt(A,(n+1)>>1),y;
		x.resize(n),y.resize(n);
		for(int i=0;i<n;i++)y[i]=A[i];
		x=Poly_add(NTT(Poly_inv(x),y),x);
		int inv2=qpow(2,mod-2);
		for(int i=0;i<n;i++)
			x[i]=mul(x[i],inv2);
		x.resize(n);
		return x;
	}
	Poly sqrt(Poly A){
		return sqrt(A,A.size());
	}
	Poly rever(Poly A){
		reverse(A.begin(),A.end());
		return A;
	}
	void div(Poly A,Poly B,Poly &C,Poly &D){
		int n=A.size(),m=B.size();
		Poly ra=rever(A),rb=rever(B);
		ra.resize(n-m+1),rb.resize(n-m+1);
		C=NTT(ra,Poly_inv(rb));
		C.resize(n-m+1);
		C=rever(C);
		D=Poly_sub(A,NTT(B,C));
		D.resize(m);
	}
}
using namespace Template_Poly;
typedef Poly poly;
int n,m;
int a[N],fac[N],ifac[N];
inline void init(int n = 100010){
	fac[0]=ifac[0]=1;for(int i=1;i<=n;i++)fac[i] = mul(fac[i-1],i);
	ifac[n]=qpow(fac[n],mod-2);for(int i=n-1;i;i--)ifac[i] = mul(ifac[i+1],i+1);
}

poly solve(int l,int r){
	if(l==r){
		poly ret(2,1);
		ret[1]=sub(0,a[l]);
		return ret;
	}
	int mid=(l+r)>>1;
	return NTT(solve(l,mid),solve(mid+1,r));
}
poly func;
inline void calc(){
	func = solve(1,n);
	func = ln(func);
	// for(int i=0;i<func.size();i++)cout << func[i] << " ";puts("");
	func = Deri(func);
	func.push_back(0);
	for(int i=(int)func.size()-1;i;i--)func[i] = sub(0,func[i-1]);
	func[0]=n;
	// for(int i=0;i<func.size();i++)cout << func[i] << " ";puts("");
}
poly A,B,C;
poly ret;
void Doit(){
	A.resize(n-1),B.resize(n-1);
	for(int i=0;i<=n-2;i++){
		A[i] = mul(qpow(i+1, m),ifac[i]);
		B[i] = mul(qpow(i+1, 2*m),ifac[i]);
	}
	//A[0] = 1
	// B = Poly_inv(B);
	C = NTT(B,Poly_inv(A));
	C.resize(n-1);
	A = ln(A);
	for(int i=0;i<=n-2;i++){
		A[i] = mul(A[i], func[i]);
		C[i] = mul(C[i], func[i]);
	}
	A = exp(A);
	ret = NTT(A,C);
	int ans=ret[n-2];
	for(int i=1;i<=n;i++)ans=mul(ans,a[i]);
	ans = mul(ans, fac[n-2]);
	printf("%d\n",ans);
}

int main()
{
	init();
	cin >> n >> m;
	for(int i=1;i<=n;i++)scanf("%d",&a[i]);
	calc();
	Doit();
}