【洛谷4002】[清华集训2017] 生成树计数（prufer序列+大力多项式）

点此看题面

• 给定一张森林，其中有(n)棵树，第(i)棵树中有(a_i)个点。
• 要求再连上(n-1)条边使得整张图成为一棵树，假设一种方案中第(i)棵树共连出(d_i)条边，则这种方案的价值为((prod_{i=1}^nd_i^m)(sum_{i=1}^nd_i^m))。
• 求所有方案的价值之和。
• (nle3 imes10^4,mle30)

(prufer)序列

不妨改为用每个点的度数-1表示(d_i)。

根据(prufer)序列的结论，这种情况下的生成树种数应该是(frac{(n-2)!}{prod_{i=1}^nd_i!})（其实就是个可重全排列）。

再考虑上每个连通块的每条边连接的点都有(a_i)种选择，因此列出答案的计算式为：

[sum_{sum_{i=1}^nd_i=n-2}frac{(n-2)!}{prod_{i=1}^nd_i!}prod_{i=1}^na_i^{d_i+1}(d_i+1)^msum_{i=1}^m(d_i+1)^m ]

稍微变个形便可以得到：

[(n-2)!prod_{i=1}^na_icdotsum_{sum_{i=1}^nd_i=n-2}prod_{i=1}^nfrac{a_i^{d_i}}{d_i!}(d_i+1)^msum_{i=1}^m(d_i+1)^m ]

我们可以最后再给答案乘上((n-2)!prod_{i=1}^na_i)这个常数，那就只需考虑其后这坨式子。

式子化简+生成函数

这个东西看起来长得非常恶心，但实际上我们可以直接把前面(prod)中的式子乘到后面的(sum)里面，将它写成：

[sum_{sum_{i=1}^nd_i=n-2}sum_{i=1}^nfrac{a_i^{d_i}}{d_i!}(d_i+1)^{2m}prod_{j ot=i}frac{a_j^{d_j}}{d_j!}(d_j+1)^m ]

这时候的式子总算长成了一个卷积的形式，我们构造出两个关于(d_i)的生成函数：

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

由于原式的(a_i^{d_i})一项中(d_i)充当指数，因此我们需要把(a_i)写到括号里面，就可以转化得到：

[[x^{n-2}]sum_{i=1}^nA(a_ix)prod_{j ot=i}B(a_jx) ]

后面式子中(j ot=i)一看就很麻烦，但又很容易消掉，只要在后面的式子中乘上它并在前面的式子中除以它即可，而这样一来的好处就是(i)与(j)独立了：

[[x^{n-2}]sum_{i=1}^nfrac AB(a_ix)prod_{j=1}^nB(a_jx) ]

其中的(prod_{j=1}^nB(a_jx))是一堆多项式乘法想想都不太可行，因此考虑写成(exp(sum_{j=1}^nln B(a_jx)))化乘为加：

[[x^{n-2}]sum_{i=1}^nfrac AB(a_ix)exp(sum_{j=1}^nln B(a_jx)) ]

其中的(frac AB)直接多项式求逆后多项式乘法，而(ln B)也就是一个多项式(ln)板子。

现在的问题在于这两部分式子前面都有一个(sum_{i=1}^n)，且括号里面都有一个(a_i)，也就是说我们需要给这两个多项式的每个第(k)次项乘上一个(sum_{i=1}^na_i^k)。

所以问题来了，如何对于所有的(k=0sim n)，快速求出(sum_{i=1}^na_i^k)。

快速求整个数列全部(k)次方和

又是一大难点。

不难想到构造这个的生成函数：

[F(x)=sum_{k}sum_{i=1}^na_i^kx^k ]

考虑到(ln(f(x))'=frac{f'(x)}{f(x)})，所以说就有：

[G(x)=ln(prod_{i=1}^n(1-a_ix))'=sum_{i=1}^nln(1-a_ix)'=sum_{i=1}^nfrac{-a_i}{1-a_ix} ]

而众所周知(frac1{1-a_ix}=sum_ka_i^kx^k)，可以继续化简得到：

[G(x)=sum_{i=1}^n-a_isum_k(a_ix)^k=-sum_ksum_{i=1}^na_i^{k+1}x^k ]

发现(G(x))和(F(x))长得很像（废话，本来就是专门为它构造的），实际上我们只要把(sum)前的系数改正、给所有(x^k)次数加(1)并添上常数项(n)，就能发现它们之间存在的关系式：

[F(x)=-xG(x)+n ]

而要求(G(x))，显然只要分治(NTT)求出(prod_{i=1}^n(1-a_ix))，然后多项式(ln)+求导即可。（所以说是怎么想到这种构造的啊！）

求出(F(x))之后，它的第(k)项系数就是(sum_{i=1}^na_i^k)。

利用它就可以求出前面多项式每一项的真正系数，再把(ln B)给(exp)回去卷上(frac AB)，这道题就终于落下帷幕了。

代码：(O(nlog^2n))

#include<bits/stdc++.h>
#define Tp template<typename Ty>
#define Ts template<typename Ty,typename... Ar>
#define Reg register
#define RI Reg int
#define Con const
#define CI Con int&
#define I inline
#define W while
#define N 30000
#define X 998244353
using namespace std;
int n,m,a[N+5],Fac[N+5],IFac[N+5],A[N+5],B[N+5],B_[N+5],S[N+5],F[N+5],ans[N+5];
I int QP(RI x,RI y) {RI t=1;W(y) y&1&&(t=1LL*t*x%X),x=1LL*x*x%X,y>>=1;return t;}
namespace Poly//多项式板子
{
	#define PR 3
	#define Init(n) P=1,L=0;W(P<=2*(n)) P<<=1,++L;
		for(i=0;i^P;++i) A[i]=B[i]=0,R[i]=((R[i>>1]>>1)|((i&1)<<L-1));
	int P,L,R[N<<2],p[N+5],A[N<<2],B[N<<2];
	I void NTT(int* s,CI op)
	{
		RI i,j,k,x,y,U,S;for(i=0;i^P;++i) i<R[i]&&(x=s[i],s[i]=s[R[i]],s[R[i]]=x);
		for(i=1;i^P;i<<=1) for(U=QP(QP(PR,op),(X-1)/(i<<1)),j=0;j^P;j+=i<<1) for(S=1,k=0;
			k^i;++k,S=1LL*S*U%X) s[j+k]=((x=s[j+k])+(y=1LL*S*s[i+j+k]%X))%X,s[i+j+k]=(x-y+X)%X;
	}
	I void Mul(CI n,int* a,int* b)//多项式乘法
	{
		RI i;Init(n);for(i=0;i<=n;++i) A[i]=a[i],B[i]=b[i];
		for(NTT(A,1),NTT(B,1),i=0;i^P;++i) A[i]=1LL*A[i]*B[i]%X;
		RI t=QP(P,X-2);for(NTT(A,X-2),i=0;i<=n;++i) a[i]=1LL*A[i]*t%X; 
	}
	I void Inv(CI n,int* a,int* b)//多项式求逆
	{
		if(!n) return (void)(b[0]=QP(a[0],X-2));RI i;Inv(n>>1,a,b);
		Init(n);for(i=0;i<=n;++i) A[i]=a[i],B[i]=b[i];
		for(NTT(A,1),NTT(B,1),i=0;i^P;++i) A[i]=(2LL*B[i]-1LL*A[i]*B[i]%X*B[i]%X+X)%X;
		RI t=QP(P,X-2);for(NTT(A,X-2),i=0;i<=n;++i) b[i]=1LL*A[i]*t%X;
	}
	I void Ln(CI n,int* a,int* b)//多项式ln
	{
		RI i;for(i=0;i<=n;++i) b[i]=0;Inv(n,a,b);
		Init(n-1);for(i=0;i<=n-1;++i) A[i]=1LL*a[i+1]*(i+1)%X,B[i]=b[i];
		for(NTT(A,1),NTT(B,1),i=0;i^P;++i) A[i]=1LL*A[i]*B[i]%X;
		RI t=QP(P,X-2);for(NTT(A,X-2),b[0]=i=0;i^n;++i) b[i+1]=1LL*A[i]*t%X*QP(i+1,X-2)%X;
	}
	int q[N+5];I void Exp(CI n,int* a,int* b)//多项式exp
	{
		if(!n) return (void)(b[0]=1);RI i;Exp(n>>1,a,b);
		Ln(n,b,p),Init(n);for(i=0;i<=n;++i) A[i]=b[i],B[i]=(!i-p[i]+a[i]+X)%X;
		for(NTT(A,1),NTT(B,1),i=0;i^P;++i) A[i]=1LL*A[i]*B[i]%X;
		RI t=QP(P,X-2);for(NTT(A,X-2),i=0;i<=n;++i) b[i]=1LL*A[i]*t%X;
	}
	int ct,f[50][N+5];I void Solve(CI l,CI r,CI rt)//分治NTT求∏(1-ax)
	{
		if(l==r) return (void)(f[rt][0]=1,f[rt][1]=X-a[l]);
		RI i,mid=l+r>>1,lc=++ct,rc=++ct;Solve(l,mid,lc),Solve(mid+1,r,rc);
		P=1,L=0;W(P<=2*(r-l+1)) P<<=1,++L;for(i=0;i^P;++i) A[i]=B[i]=0,R[i]=(R[i>>1]>>1)|((i&1)<<L-1);
		for(i=0;i<=mid-l+1;++i) A[i]=f[lc][i];for(i=0;i<=r-mid;++i) B[i]=f[rc][i];
		for(NTT(A,1),NTT(B,1),i=0;i^P;++i) A[i]=1LL*A[i]*B[i]%X;
		RI t=QP(P,X-2);for(NTT(A,X-2),i=0;i<=r-l+1;++i) f[rt][i]=1LL*A[i]*t%X;ct-=2;
	}
}
int main()
{
	RI i;for(scanf("%d%d",&n,&m),i=1;i<=n;++i) scanf("%d",a+i);
	for(Fac[0]=i=1;i<=n;++i) Fac[i]=1LL*Fac[i-1]*i%X;//预处理阶乘
	for(IFac[i=n]=QP(Fac[n],X-2);i;--i) IFac[i-1]=1LL*IFac[i]*i%X;//预处理阶乘逆元
	for(i=0;i<=n;++i) A[i]=1LL*QP(i+1,2*m)*IFac[i]%X,B[i]=1LL*QP(i+1,m)*IFac[i]%X;//求出初始的A,B
	Poly::Inv(n,B,B_),Poly::Mul(n,A,B_),Poly::Ln(n,B,S),Poly::Solve(1,n,0),Poly::Ln(n,Poly::f[0],F);//A=A/B,S=lnB,F=ln∏(1-ax)
	for(i=0;i<=n;++i) F[i]=1LL*F[i+1]*(i+1)%X;for(i=n;i;--i) F[i]=(X-1LL)*F[i-1]%X;F[0]=n;//对F求导，然后变成-xF+n成为真正需要的F
	for(i=0;i<=n;++i) A[i]=1LL*F[i]*A[i]%X,S[i]=1LL*F[i]*S[i]%X;//给两个多项式的第i项乘上∑(a^i)
	Poly::Exp(n,S,ans),Poly::Mul(n,ans,A);RI t=1LL*ans[n-2]*Fac[n-2]%X;//把S给exp回去再和A卷起来，t初始化为答案多项式的第n-2次项乘上常数中的(n-2)!
	for(i=1;i<=n;++i) t=1LL*t*a[i]%X;return printf("%d
",t),0;//乘上常数中的∏a
}