LOJ2320 生成树计数

多项式变量添系数求和

给一个 (t) 次多项式 (P) 和 (n) 个数 (a_1sim a_n)，求 (sum_{i=1}^n P(a_ix))。

题解

若 (P=sum_{k=0}^t p_kx^k) 那么

[sum_{i=1}^n P(a_ix)=sum_{k=0}^t p_ksum_{i=1}^n a_i^kx^k ]

先求 (A) 的 (0sim t) 次等指数幂和 (s_A(x)=sum_{k=0}^infty x^k sum_{i=1}^na_i^k)，然后与 (P) 的系数对位相乘即可。

等指数幂和用

[s_A(x)=sum_{i=1}^nfrac{1}{1-a_ix}=frac{sum_{i=1}^nprod_{j e i}(1-a_jx)}{prod_{i=1}^n(1-a_ix)} ]

求出，时间复杂度 (O(nlog^2n+tlog t))。

考虑 (sum_{i=1}^nprod_{j e i}(1-a_jx)) 与 (prod_{i=1}^n(1-a_ix)) 的关系。可以发现后者任意 (i) 次项在前者中恰好被算了 (n-i) 次。所以用分治NTT求出后者后简单处理即可得到求前者。

多项式变量添系数求积

给一个 (t) 次多项式 (P) 和 (n) 个数 (a_1sim a_n)，求 (prod_{i=1}^n P(a_ix))。

题解

套上 (ln) 将积转为和

[prod_{i=1}^nP(a_ix)=expleft(sum_{i=1}^nln P(a_ix) ight)\ =expleft(sum_{i=1}^n(ln P)(a_i x) ight) ]

因此我们对 (P) 取 (ln)，执行一遍“多项式变量添系数求和”，再 (exp) 回来即可。时间复杂度 (O(nlog^2n+tlog t))。

生成树计数

在一个 (s) 个点的图中，存在 (s-n) 条边，使图中形成了 (n) 个连通块，第 (i) 个连通块中有 (a_i) 个点。

现在我们需要再连接 (n-1) 条边，使该图变成一棵树。对一种连边方案，设原图中第 (i) 个连通块连出了 (d_i) 条边，那么这棵树 (T) 的价值为：

[mathrm{val}(T) = left(prod_{i=1}^{n} d_i^m ight)left(sum_{i=1}^{n} d_i^m ight) ]

你的任务是求出所有可能的生成树的价值之和，对 (998244353) 取模。

(100\%) 的数据中，(n le 3 imes 10^4,m le 30)。

魔幻做法

https://rqy.moe/Solutions/bzoj5119/

连加的性质显然比连乘多，所以我们把 (mathrm{val}(T)) 定义式中的连加提到前面来。

[mathrm{val}(T)=sum_{i=1}^nd_i^mprod_{j=1}^{n} d_j^m\ =sum_{i=1}^nd_i^{2m}prod_{j eq i} d_j^m ]

先用枚举树的形态列出暴力求和式

[ans=sum_Tsum_{i=1}^nd_i^{2m}prod_{j eq i}d_j^m ]

答案只与度数有关，所以我们联想到Prufer序列。我们对连通块做Prufer序列，枚举每个连通块出现的次数 (d_i)，那么连通块的连边数应为 (d_i+1)。同时由于此时的Prufer序列装的是连通块，所以我们还要对每个连通块乘以连边方案 (a_i^{d_i+1})。

[ans=sum_D frac{(n-2)!}{prod d!}prod a^{d+1}sum_{i=1}^n (d_{i}+1)^{2m} prod_{j eq i} (d_j+1)^m\ =sum_D frac{(n-2)!}{prod d!} sum_{i=1}^n a_i^{d_i+1}(d_{i}+1)^{2m} prod_{j eq i} a_j^{d_j+1}(d_j+1)^m\ =(n-2)!sum_{i=1}^nsum_D frac{a_i^{d_i+1}(d_{i}+1)^{2m}}{d_i!} prod_{j eq i} frac{a_j^{d_j+1}(d_j+1)^m}{d_j!} ]

用生成函数优化 (D) 的枚举。

[A_i(x)=sum_{k=0}^infty frac{a_i^{k+1}(k+1)^m}{k!}x^k\ B_i(x)=sum_{k=0}^infty frac{a_i^{k+1}(k+1)^{2m}}{k!}x^k\ ans=(n-2)!sum_{i=1}^n left(B_i prod_{j eq i} A_j ight)[x^{n-2}] ]

接下来我们要寻找快速计算 (B_i prod_{j eq i} A_j) 的方法。注意到 (A,B) 的形式是指数型生成函数，而系数与指数没有对上，所以不妨对 (A) 积分。

[int A_i(x)mathrm dx=sum_{k=1}^infty a_i^kk^mfrac{x^k}{k!} ]

魔幻的一步——斯特林套路。

[int A_i(x)mathrm dx=sum_{k=1}^infty a_i^kfrac{x^k}{k!}sum_{j=0}^megin{Bmatrix}m\jend{Bmatrix}k^underline{j}\ =sum_{j=0}^megin{Bmatrix}m\jend{Bmatrix}sum_{k=j}^infty a_i^kk^underline{j}frac{x^k}{k!}\ =sum_{j=0}^megin{Bmatrix}m\jend{Bmatrix}sum_{k=j}^infty a_i^kfrac{x^k}{(k-j)!} ]

注意到 (frac{1}{(k-j)!}) 的形式，我们把后面那个求和式都提一个 (k-j) 出来。

[int A_i(x)mathrm dx=sum_{j=0}^megin{Bmatrix}m\jend{Bmatrix}a_i^jx^jsum_{k=0}^infty a_i^kfrac{x^k}{k!}\ =e^{a_ix}sum_{j=0}^megin{Bmatrix}m\jend{Bmatrix}a_i^jx^j ]

(B) 与 (A) 的唯一区别就是 (k^m) 变成了 (k^{2m})，所以 (int B) 就简单变化一下斯特林套路枚举的范围即可。

[int B_i(x)mathrm dx=e^{a_ix}sum_{j=0}^{2m}egin{Bmatrix}2m\jend{Bmatrix}a_i^jx_j\ ]

现在我们再来求导

[A_i(x)=e^{a_x}sum_{j=0}^megin{Bmatrix}m\jend{Bmatrix}a_i^{j+1}x^j+e^{a_ix}sum_{j=0}^megin{Bmatrix}m\j+1end{Bmatrix}a_i^{j+1}(j+1)x^j\ =e^{a_ix}sum_{j=0}^mleft(egin{Bmatrix}m\jend{Bmatrix}+(j+1)egin{Bmatrix}m\j+1end{Bmatrix} ight)a_i^{j+1}x^j\ =e^{a_ix}sum_{j=0}^megin{Bmatrix}m+1\j+1end{Bmatrix}a_i^{j+1}x^j\ B_i(x)=e^{a_ix}sum_{j=0}^{2m}egin{Bmatrix}2m+1\j+1end{Bmatrix}a_i^{j+1}x^j ]

最后我们再来推 (B_iprod_{j eq i}A_j) 的式子。

[B_i(x)prod_{j eq i}A_i(x)=e^{sum ax}left(sum_{k=0}^{2m}egin{Bmatrix}2m+1\k+1end{Bmatrix}a_i^{k+1}x^k ight)prod_{j eq i}left(sum_{k=0}^megin{Bmatrix}m+1\k+1end{Bmatrix}a_j^{k+1}x^k ight) ]

注意到前面的 (e^{sum ax}) 是个定值，而后面的可以分治NTT求，所以这题就做完了。时间复杂度 (O(nmlog^2n))。

如何分治NTT呢？注意到我们只需要对所有的 (B_iprod_{j eq i} A_j) 求和，所以可以像solve(l,mid)求出 (sum_{i=l}^{mid}B_iprod_{j=l,j eq i}^{mid}A_j) 后让它乘以 (prod_{j=mid+1}^rA_j)这样计算。

套路优化

提取 (B_i prod_{j eq i} A_j) 的公因子 (prod a)，那么

[A(x)=e^xsum_{k=0}^megin{Bmatrix}m+1\k+1end{Bmatrix}x^k\ B(x)=e^xsum_{k=0}^megin{Bmatrix}2m+1\k+1end{Bmatrix}x^k\ A_i(x)=A(a_ix)\ B_i(x)=B(a_ix) ]

再对 (sum_{i=1}^nB_iprod_{j eq i}A_j) 变形，

[sum_{i=1}^nB_i(x)prod_{j eq i}A_j(x)=sum_{i=1}^nfrac{B_i(x)}{A_i(x)}prod_{j=1}^nA_j(x)\ left(sum_{i=1}^nleft(frac{B}{A} ight)(a_ix) ight)left(prod_{i=1}^nA(a_ix) ight) ]

套用“多项式变量添系数求和”与“多项式变量添系数求积”即可，这部分的复杂度是 (O(nlog^2 n))。

另外需要求出第二类斯特林数的第 (m) 行的前 (n) 项，关键的复杂度在于快速幂，(O(nlog m))。

组合做法

https://blog.csdn.net/qq_39972971/article/details/80786441

考虑给定式子的组合意义：

1. 写出一个Prufer序列，在序列后顺序补足 (1) 到 (n) 各一个，使得各数出现次数等于其树上的点度。每个数 (x) 的权值为 (a_x)。

2. 重复 (m) 次后述染色：在每一种数中选择一个位置涂色（两次染色之间独立）。

3. 选取一种数 (x)，重复 (m) 次后述染色：在写着 (x) 的位置中选择一个涂色。

记 (f_i(j)) 表示考虑了 (1) 到 (i)，尚未选取过 (x)（进行步骤3），Prufer序列中的 (n-2) 个位置有 (j) 个位置已经涂了色，可行的方案数与已经考虑过的数的权值乘积的和；(g_i(j)) 表示考虑了 (1) 到 (i)，已经选取过 (x)，Prufer序列中的 (n-2) 个位置有 (j) 个位置已经涂了色，可行的方案数与已经考虑过的数的权值乘积的和。最终计算答案有

[ans=sum_{i=0}^{n-2}g_n(i) imes left(sum a ight)^{n-2-i}。 ]

转移时枚举对于 (i)，有多少Prufer序列中的位置被涂了色，则有

[f_i(j)=sum_{k=0}^{m}f_{i-1}(j-k) imes inom{n-2-(j-k)}{k}a_i^{k+1}left(egin{Bmatrix}m\kend{Bmatrix}k!+egin{Bmatrix}m\k+1end{Bmatrix}(k+1)! ight)\ g_i(j)=sum_{k=0}^{2m}f_{i-1,j-k} imes inom{n-2-(j-k)}{k}a_i^{k+1}left(egin{Bmatrix}2m\kend{Bmatrix}k!+egin{Bmatrix}2m\k+1end{Bmatrix}(k+1)! ight)\ +sum_{k=0}^{m}g_{i-1,j-k} imes inom{n-2-(j-k)}{k}a_i^{k+1}left(egin{Bmatrix}m\kend{Bmatrix}k!+egin{Bmatrix}m\k+1end{Bmatrix}(k+1)! ight) ]

直接实现上述DP，时间复杂度 (O(n^2m))。

上述DP的式子可以的化为卷积的形式：

[F_i=F_{i-1} imes A_i\ G_i=G_{i-1} imes A_i+F_{i-1} imes B_i ]

其中 (A_i,B_i) 为转移系数多项式，(A_i) 的次数为 (m)，(B_i) 的次数为 (2m)，我们希望求出 (G_n) 系数的最后 (n-1) 项。

分治NTT，将系数卷积起来，并每一步对 (x^n) 取模即可。组合数的形式有些特殊，不过只需要按照EGF卷积，最后每项乘以 (frac{(n-2)!}{i!}) 而不是普通的 (i!) 就好了。时间复杂度 (O(nmlog^2 n))。

套路优化

注意到 (G=sum_{i=1}^n A_iprod_{j eq i}B_j)，使用“魔幻做法”同样的手段优化即可，这部分复杂度 (O(n log^2 n))。

同样需要第二类斯特林数求行，关键的复杂度在于快速幂，(O(nlog m))。

本质做法

https://loj.ac/article/1257

前半部分

先不考虑后面那个 (sum)，那么可以表示为若干个 (w(d)=d^m) 的积。

用EGF分配Prufer序列的位置，若第 (i) 个块分到 (d_i) 个位置，那么其度数为 (d_i+1)。并且要为每个度数决定连接哪个点，即 (a_i^{d_i+1})。

用EGF分配位置……生成函数本质上是工具。

[ans=(prod a)sum_Dfrac{(n-2)!}{prod d!}prod_{i=1}^n a_i^{d_i}w(d_i+1) ]

这里把公因子 (prod a) 提到了前面。(frac{(n-2)!}{prod d!}) 是分配位置，用EGF搞掉。

那么块 (i) 的EGF是 (F_i(x) = sum_{k=0}^inftyfrac{x^k}{k!} a_i^k w(k+1))（(w) 不要漏写+1）。

可以发现如果设 (F(x) = sum_{k=0}^inftyfrac{x^k}{k!} w(k+1))，那么 (F_i(x)=F(a_ix))。

把每块的EGF乘起来得到

[ans = (prod a)(n-2)!left(prod_{i=1}^n F(a_ix) ight)[x^{n-2}] ]

以上推理只要树权值满足 (prod w(d)) 的形式即可进行。

出题人给的 (d^m) 本质上是关于 (d) 的函数。

后半部分

前半部分讲到我们要求 (prod_{i=1}^n F(a_ix))，其中 (F(x) = sum_{k=0}^inftyfrac{x^k}{k!} w(k+1))。使用“多项式变量添系数求积”即可。

现在考虑原题的 (val(T) = (prod_{i=1}^n d_i^m ) (sum_{i=1}^n d_i^m))，尝试表示成积的形式，发现它就是硬点了某个 (i) 不贡献 (w(d_i)=d_i^m) 而贡献 (w'(d_i)=d_i^{2m})。

出题人右半边求和的本质就是更改某个连通块贡献。

相应地再设 (F'(x) = sumfrac{x^k}{k!} w'(k+1))，那么就要枚举 (i) 并把一个参与 (prod) 的 (F(a_ix)) 替换成 (F'(a_ix))。

[ans=(prod a)(n-2)!left(prod_{i=1}^n F(a_ix) ight)left(sum_{i=1}^nfrac{F'(a_ix)}{F(a_ix)} ight)[x^{n-2}] ]

求出 (G=frac{F'}F) 再执行一遍“多项式变量添系数求和”即可。

综上，共需求1次等幂和(含1分治FFT、1求逆、1乘法)、1 (ln)、1 (exp)、1求逆、2乘法。复杂度 (O(nlog^2n+nlog m))，其中 (log m) 是快速幂。

(m) 可以做1e9。提交记录 3e4 跑了不到 160ms，预计 2.5s 内可以做到 2.6e5 或 5.2e5。

CO int N=65536;
int omg[2][N],rev[N],inv[N];

void NTT(poly&a,int dir){
	int lim=a.size(),len=log2(lim);
	for(int i=0;i<lim;++i) rev[i]=rev[i>>1]>>1|(i&1)<<(len-1);
	for(int i=0;i<lim;++i)if(i<rev[i]) swap(a[i],a[rev[i]]);
	for(int i=1;i<lim;i<<=1)
		for(int j=0;j<lim;j+=i<<1)for(int k=0;k<i;++k){
			int t=mul(omg[dir][N/(i<<1)*k],a[j+i+k]);
			a[j+i+k]=add(a[j+k],mod-t),a[j+k]=add(a[j+k],t);
		}
	if(dir==1){
		int ilim=fpow(lim,mod-2);
		for(int i=0;i<lim;++i) a[i]=mul(a[i],ilim);
	}
}
poly operator~(poly a){
	int n=a.size();
	poly b(1,fpow(a[0],mod-2));
	if(n==1) return b;
	int lim=2;
	for(;lim<n;lim<<=1){
		poly a1(a.begin(),a.begin()+lim);
		a1.resize(lim<<1),NTT(a1,0);
		b.resize(lim<<1),NTT(b,0);
		for(int i=0;i<lim<<1;++i) b[i]=mul(2+mod-mul(a1[i],b[i]),b[i]);
		NTT(b,1),b.resize(lim);
	}
	a.resize(lim<<1),NTT(a,0);
	b.resize(lim<<1),NTT(b,0);
	for(int i=0;i<lim<<1;++i) b[i]=mul(2+mod-mul(a[i],b[i]),b[i]);
	NTT(b,1),b.resize(n);
	return b;
}
poly log(poly a){
	int n=a.size();
	poly b=~a;
	for(int i=0;i<n-1;++i) a[i]=mul(a[i+1],i+1);
	a.resize(n-1);
	int lim=1<<(int)ceil(log2(2*n-2));
	a.resize(lim),NTT(a,0);
	b.resize(lim),NTT(b,0);
	for(int i=0;i<lim;++i) a[i]=mul(a[i],b[i]);
	NTT(a,1),a.resize(n);
	for(int i=n-1;i>=1;--i) a[i]=mul(a[i-1],inv[i]);
	return a[0]=0,a;
}
poly exp(poly a){
	int n=a.size();
	poly b(1,1); // a[0]=0
	if(n==1) return b;
	int lim=2;
	for(;lim<n;lim<<=1){
		poly a1(a.begin(),a.begin()+lim);
		b.resize(lim);poly b1=log(b);
		a1.resize(lim<<1),NTT(a1,0);
		b.resize(lim<<1),NTT(b,0);
		b1.resize(lim<<1),NTT(b1,0);
		for(int i=0;i<lim<<1;++i) b[i]=mul(add(1+a1[i],mod-b1[i]),b[i]);
		NTT(b,1),b.resize(lim);
	}
	b.resize(lim);poly b1=log(b);
	a.resize(lim<<1),NTT(a,0);
	b.resize(lim<<1),NTT(b,0);
	b1.resize(lim<<1),NTT(b1,0);
	for(int i=0;i<lim<<1;++i) b[i]=mul(add(1+a[i],mod-b1[i]),b[i]);
	NTT(b,1),b.resize(n);
	return b;
}
poly operator*(poly a,poly b){
	int n=a.size()-1,m=b.size()-1;
	int lim=1<<(int)ceil(log2(n+m+1));
	a.resize(lim),NTT(a,0);
	b.resize(lim),NTT(b,0);
	for(int i=0;i<lim;++i) a[i]=mul(a[i],b[i]);
	NTT(a,1),a.resize(n+m+1);
	return a;
}

int a[N];

poly prod(int l,int r){
	if(l==r) return (poly){1,mod-a[l]};
	int mid=(l+r)>>1;
	return prod(l,mid)*prod(mid+1,r);
}

int main(){
	omg[0][0]=1,omg[0][1]=fpow(3,(mod-1)/N);
	omg[1][0]=1,omg[1][1]=fpow(omg[0][1],mod-2);
	for(int i=2;i<N;++i){
		omg[0][i]=mul(omg[0][i-1],omg[0][1]);
		omg[1][i]=mul(omg[1][i-1],omg[1][1]);
	}
	inv[0]=inv[1]=1;
	for(int i=2;i<N;++i) inv[i]=mul(mod-mod/i,inv[mod%i]);
	
	int n=read<int>(),m=read<int>();
	if(n==1){
		puts(m==0?"1":"0");
		return 0;
	}
	poly F(n-1),G(n-1);
	int ifac=1;
	for(int i=0;i<n-1;++i){
		int t=fpow(i+1,m);
		ifac=mul(ifac,inv[i]);
		F[i]=mul(t,ifac);
		G[i]=mul(t,F[i]);
	}
	G=G*~F,G.resize(n-1);
	F=log(F);
	
	int ans=fpow(ifac,mod-2);
	for(int i=1;i<=n;++i) read(a[i]),ans=mul(ans,a[i]);
	poly Q=prod(1,n),P(n-1);
	Q.resize(n-1);
	for(int i=0;i<n-1;++i) P[i]=mul(Q[i],n-i);
	P=P*~Q,P.resize(n-1);
	
	for(int i=0;i<n-1;++i){
		F[i]=mul(F[i],P[i]);
		G[i]=mul(G[i],P[i]);
	}
	ans=mul(ans,(exp(F)*G)[n-2]);
	printf("%d
",ans);
	return 0;
}