牛客周赛 Round50 E-小红的树上移动（期望dp+逆元）

onlypasserby / 2024-07-10 / 原文

E-小红的树上移动

题目：

题意：在一个树上从根节点移动，每次都会向更深的下一层走，如果此时已经是叶子节点没有下一层就会停留在这里。求出移动次数的期望，移动次数就是从根节点1开始到此节点的深度。

思路：

画一个草图不难看出其实在同一层中，到达每个点的概率是一样的。并且，对于每一层的节点概率就是这一层节点个数乘能到达这一层的概率。其中能到达这一层的概率就是到达上一层的概率乘上一层的（非叶子点数/总点数）。

那么，总结一下我们能得到如此的递推关系dp[i]=dp[i-1]*(a-b)/a

dp[i]为到达i+1层的概率，a为总点数，b为叶子点数

那么我们要计算的答案期望就是dp[i-1]乘(b/a)乘d[i],d[i]就是这一层的深度

以下是代码，！注意求逆元时的取模问题

点击查看代码

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define endl "\n"
#define ios ios::sync_with_stdio(0),cin.tie(0),cout.tie(0)
const int N = 2e6+10;
const int p = 998244353;
int f[N],inv[N],yz[N],dep[N];
int km(int a,int b)
{
	int res=1;
	while(b)
	{
		if(b&1) res=res*a%p;
		a=a*a%p;
		b>>=1;
	}
	return res;
}
int n,d[N],dp[N],mx=0;
vector<int> h[N];
bool vis[N];
void add(int a,int b)
{
	h[a].push_back(b);
	h[b].push_back(a);
} 
void bfs()
{
	queue<int> q;
	q.push(1);
	vis[1]=1,d[1]=0;
	while(q.size())
	{
		int u=q.front();
		q.pop();
		for(auto v:h[u])
		{
			if(vis[v]) continue;
			vis[v]=1;
			d[v]=d[u]+1;//算出深度 
			dep[d[v]]++;//纪录每个深度点的个数 
            if(h[v].size()==1&&v!=1) yz[d[v]]++;//纪录叶子节点个数 
			mx=max(mx,d[v]);//纪录下最大深度 
			q.push(v);
		}
	}
}
void solve()
{
	cin>>n;
	for(int i=1;i<n;i++) {
		int u,v;
		cin>>u>>v;
		add(u,v);
	}
	bfs();
	dp[0]=1;
	int ans=0;
    for(int i=1;i<=mx;i++)//递推求答案过程,注意求逆元时取模 
    {
        int a=dep[i],b=yz[i];
        ans+=i%p*dp[i-1]%p*b%p*km(a,p-2)%p;
        ans%=p;
        dp[i]=dp[i-1]%p*(a-b)%p*km(a,p-2)%p;
        dp[i]%=p;
    }
	cout<<ans<<endl;
}
signed main()
{
	ios;
	int t=1;
	//cin>>t;
	while(t--)
	{
		solve() ;
	}
	return 0;
}