費馬小定理-期望dp

haihaichibaola發表於2024-07-10

E - 小紅的樹上移動 (Nowcoder85687 E)

題目大意

小紅有一棵 𝑛 個點的樹，根節點為 1，有一個物塊在根節點上，每次它會等機率隨機移動到當前節點的其中一個子節點，而後等機率隨機傳送到一個同深度節點，若此時它位於葉子節點，則停止移動。
求其移動到子節點的次數的期望值，答案對 998244353 取模。

解題思路

我們可以分別對每一層的情況進行討論。首先使用dfs求出樹在每一個深度時具有的節點數量和葉子節點數量。然後根據題意，我們可以發現，每一層達到葉子節點的期望可以表示為，上一層到達下一層的機率這一層中葉子節點的機率再遍歷次數（樹的深度），即(\(cnt/sum\times p\times\ deep)\)%$mod $然後將每一層的期望累加皆可。

需要注意的是，在計算除法的時候，也需要對結果進行取餘數。但是會遇到一些問題。

常見的取餘：

  • (a + b) % p = (a%p + b%p) %p
  • (a - b) % p = ((a%p - b%p) + p) %p
  • (a * b) % p = (a%p)*(b%p) %p

當我們要計算(\(a\times b)\)%$mod $時，較大的資料往往會超出資料範圍。這時使用快速乘法可以解決這個問題。

typedef long long LL;
LL q_mul(LL a, LL b, LL p){//快速乘法取模
    LL ans = 0;
    while (b){
        if(b&1LL) ans=(ans+a)%p;
        //or ans=(ans+(b%2*a)%p)%p;
        a = (a +a) % p;
        b >>= 1;
    }
    return ans;
}

對於除法取餘，常用的取餘方式是不成立的，所以要使用費馬小定理（不做證明）。

  • a / b=a * pow_mod(b,mod-2,mod)
int mi(int a,int b,int MOD){
    int res = 1;
    while(b){
        if(b&1)res=res*a%MOD;
        b>>=1;
        a=a*a%MOD;
    }
    return res%MOD;
}
////////////////////////////////////////////////////
int main(){
    int op=((cnt[deep]%mod)*(mi(sum[deep],mod-2,mod))%mod);
    //(cnt[deep]/sum[deep])%mod
}

然後使用乘法取餘即可。

對於這道題，我們就可以使用費馬小定理解決上述的問題。

#pragma GCC optimize("O2")
#pragma GCC optimize("O3")
#include <bits/stdc++.h>
#define int long long
using namespace std;
const int N = 1000500;
int h[N*2],e[N*2],ne[N*2],idx;
int n,max_deep,cnt[N],sum[N];
int mod=998244353; 
int ans;
int mi(int a,int b,int MOD){
	int res = 1;
	while(b){
		if(b&1)res=res*a%MOD;
		b>>=1;
		a=a*a%MOD;
	}
	return res%MOD;
}
void add(int a,int b){
	e[idx]=b,ne[idx]=h[a],h[a]=idx++;
}
void dfs1(int u,int fu,int deep){
	//避免訪問到父節點，更新最大深度
	max_deep=max(max_deep,deep);
	bool flag=false;
	for(int i=h[u];i!=-1;i=ne[i]){
		int j=e[i];
		if(j==fu)continue;
		flag=true;
		dfs1(j,u,deep+1);
	} 
	if(!flag){
		cnt[deep]++;
	}
	sum[deep]++;
}
//能夠到達下一層的機率p：上一層非葉子節點的比例 
//對於每一層，q：葉子節點的機率為這一層葉子節點的比例
//每一層落入葉子結點的期望為q*p*deep，累加得到整棵樹上到達葉子節點的期望
void dfs2(int deep,int zi,int mu){
	//zi：到達上一層所有非葉子節點的情況
	//mu：到達上一層所有節點的情況 
	//zi/mu就是到達當前層的機率
	//計算完成後更新當前層的zi和mu，進入下一層即可
	if(deep>max_deep)return;
	if(cnt[deep]){
		//有葉子節點的情況 
		int op=((cnt[deep]%mod)*(mi(sum[deep],mod-2,mod))%mod);
		op=(op*(zi%mod))%mod*mi(mu,mod-2,mod)%mod;
		op=op*deep%mod%mod;
		ans=(ans+op)%mod;
	}
	dfs2(deep+1,zi*(sum[deep]-cnt[deep]+mod)%mod,mu*sum[deep]%mod);
} 
signed main(){
	scanf("%d",&n);
	memset(h,-1,sizeof(h));
	for(int i=1;i<n;i++){
		int u,v;
		scanf("%d %d",&u,&v);
		add(u,v);
		add(v,u);
	}
	dfs1(1,-1,0);
	dfs2(0,1,1);
	cout<<ans%mod<<endl;	
	return 0;
}

於是我們使用了兩次dfs和部分最佳化解決了這個問題。

對於這個問題，我們也可以使用期望dp去解決。

對於每個節點i，我們可以發現i節點能夠向下轉移（即到達i+1節點的機率）\(dp[i]=dp[i-1]*(a-b)/a\)

a為全部節點的個數，b為葉子節點的個數

則轉移到i層停止的期望次數就是\(dp[i-1]*a/b*i\)

#pragma GCC optimize("O2")
#pragma GCC optimize("O3")
#include <bits/stdc++.h>
#define int long long
using namespace std;
const int N = 1000500;
int h[N*2],e[N*2],ne[N*2],idx;
int n,max_deep,cnt[N],sum[N],dp[N]; 
int mod=998244353; 
int ans;
int mi(int a,int b,int MOD){
	int res = 1;
	while(b){
		if(b&1)res=res*a%MOD;
		b>>=1;
		a=a*a%MOD;
	}
	return res%MOD;
}
void add(int a,int b){
	e[idx]=b,ne[idx]=h[a],h[a]=idx++;
}
void dfs1(int u,int fu,int deep){
	//避免訪問到父節點，更新最大深度
	max_deep=max(max_deep,deep);
	bool flag=false;
	for(int i=h[u];i!=-1;i=ne[i]){
		int j=e[i];
		if(j==fu)continue;
		flag=true;
		dfs1(j,u,deep+1);
	} 
	if(!flag){
		cnt[deep]++;
	}
	sum[deep]++;
}
//對於每一層，到達這一層的機率為 dp[i]=dp[i-1]*(a-b)/a
//dp[i]為到第i層能夠向下轉移的機率，a為總點數，b為葉子點數
void solve(){
	dp[0]=1;//根節點一定可以向下轉移
	for(int i=1;i<=max_deep;i++){
		int a=sum[i],b=cnt[i];
		ans+=dp[i-1]%mod*b%mod*mi(a,mod-2,mod)%mod*i%mod;
		ans%=mod;
		dp[i]=dp[i-1]%mod*(a-b+mod)%mod*mi(a,mod-2,mod)%mod;
		dp[i]%=mod;
	} 
}
signed main(){
	scanf("%d",&n);
	memset(h,-1,sizeof(h));
	for(int i=1;i<n;i++){
		int u,v;
		scanf("%d %d",&u,&v);
		add(u,v);
		add(v,u);
	}
	dfs1(1,-1,0);
	solve();
	cout<<ans%mod<<endl;	
	return 0;
}

相關文章