2021湖南多校对抗赛第四场 I - Tree

题目链接：

https://vjudge.net/contest/432769#problem/I

Solution：

换根dp

令\(dp[u][t]\)表示\(u\)下面所有子孙节点到它的路径长度的t次方的和，现考虑如何转移

设\(u\)的父节点为\(x\)，\(u\)到\(x\)的树边边权为\(w\)。假定\(u\)下面有m个节点，这m个节点到\(u\)的路径长度分别为\(d_1,d_2,...,d_m\)，有\(dp[u][t]=\sum\limits_{i=1}^md_i^t\)

那么，节点\(u\)对其父亲\(x\)的贡献就是\(u\)下面每个节点到\(u\)的路径长度加上\(w\)以及\(w\)这条边本身。也就是：

\[\sum\limits_{i=1}^m(w+d_i)^t+w^t \]

将这个式子二项式展开，也就是：

\[\sum\limits_{i=1}^m(\sum\limits_{j=0}^tC_t^jw^jd_i^{t-j}) \]

将二项式系数及\(w\)提出来，整理式子，得:

\[\sum\limits_{j=0}^{t}(C_t^{j}w^j\sum\limits_{i=1}^md_i^{t-j}) \]

注意到这个式子可以用dp数组代换，可化为：

\[\sum\limits_{j=0}^{t}(C_t^{j}w^jdp[u][t-j]) \]

也就是说，每个节点的dp值可由其儿子节点转移上来，转移系数就是二项式系数，我们这样就可以做出来每个点的dp数组了

然而这个dp数组只包含了每个节点它下面的节点到它的路径，我们考虑如何求出包含所有点的答案

令\(f[u][t]\)表示所有节点到\(u\)的路径的t次方之和，可以用一个类似换根的思想，同样通过二项式系数转移求得

令\(r[u][t]\)表示除开以\(u\)为根的子树后，剩余点到\(u\)的父亲节点\(x\)的路径t次方之和,推导思路类似，只不过这次我们由上往下进行二项式转移，不难推出状态转移方程：

\[r[u][t]=f[x][t]-\sum\limits_{i=0}^tC_t^iw^idp[u][t-i] \]

\[f[u][t]=dp[u][t]+\sum\limits_{i=0}^tC_t^iw^ir[u][t-i] \]

处理出了f数组后，就可以算期望了

代码如下：

#include<bits/stdc++.h>
using namespace std;
const long long M=998244353;
struct front_star{
	int to,next;
	long long w;
}e[200005];
int n,k,cnt=0;
int head[100005];
long long dp[100005][16],f[100005][16],r[100005][16],C[16][16];
long long fp(long long base,long long power) {
    long long result = 1;
    while (power > 0) {
        if (power & 1) {
            result =( (result%M) * (base % M))%M;
        }
        power >>= 1;
        base =( (base%M )* (base%M) )% M;
    }
    return result;
}
void addedge(int u,int v,long long val)
{
	cnt++;
	e[cnt].w=val;
	e[cnt].to=v;
	e[cnt].next=head[u];
	head[u]=cnt;
}
void fdfs(int u,int fa,int t)
{
	for(int i=head[u];~i;i=e[i].next)
	{
		int v=e[i].to;
		if(v!=fa)
		{
			fdfs(v,u,t);
			if(t==0)
			   dp[u][t]+=dp[v][t]+1;
			else
			{
				long long ew=1;
				for(int j=0;j<=t;j++)
				{
					 dp[u][t]=((dp[u][t]%M)+((((C[t][j]%M)*(ew%M))%M)*(dp[v][t-j]%M))%M)%M;
					 if(j!=t)
					    ew=((ew%M)*(e[i].w%M))%M;
				}
				dp[u][t]=(dp[u][t]%M+ew%M)%M;
			}   
		} 
	}
}
void sdfs(int u,int fa,int t,int ec)
{
	if(u==1)
	   f[u][t]=dp[u][t];
	else
	{
		long long ew=1;
		r[u][t]=f[fa][t];
		for(int i=0;i<=t;i++)
		{
			r[u][t]=((r[u][t]%M)-((((C[t][i]%M)*(ew%M))%M)*(dp[u][t-i]%M))%M+M)%M;
			if(i!=t)
			   ew=((ew%M)*(e[ec].w%M))%M;
		}
		r[u][t]=(r[u][t]%M-ew%M+M)%M;
		ew=1;
		f[u][t]=dp[u][t];
		for(int i=0;i<=t;i++)
		{
			f[u][t]=((f[u][t]%M)+((((C[t][i]%M)*(ew%M))%M)*(r[u][t-i]%M))%M)%M;
			if(i!=t)
			   ew=((ew%M)*(e[ec].w%M))%M;
		}
		f[u][t]=(f[u][t]%M+ew%M)%M;
	}
	for(int i=head[u];~i;i=e[i].next)
	{
		int v=e[i].to;
		if(v!=fa)
		   sdfs(v,u,t,i);
	}   
}
int main()
{
	scanf("%d%d",&n,&k);
	memset(head,-1,sizeof(head));
	memset(dp,0,sizeof(dp));
	for(int i=1;i<=n-1;i++)
	{
		int a,b;
		long long c;
		scanf("%d%d%lld",&a,&b,&c);
		addedge(a,b,c);
		addedge(b,a,c);
	}
	C[0][0]=1;
	for(int i=1;i<=k;i++)
	{
		C[i][0]=1;
	    for(int j=1;j<=i;j++)
	    {
	    	C[i][j]=((C[i][j-1]%M)*((i-j+1)%M))%M;
	    	C[i][j]=((C[i][j]%M)*fp(j,M-2))%M;
		}
    }
    for(int i=0;i<=k;i++)
    {
    	fdfs(1,0,i);
    	sdfs(1,0,i,0);
	}
	long long p=n,ans=0;
	long long bs=p*p%M;
	for(int i=1;i<=n;i++)
		ans=(ans%M+f[i][k]%M)%M;
	ans=((ans%M)*(fp(bs,M-2)%M))%M;
	printf("%lld\n",ans);		
	return 0;
}