#树形dp#BZOJ 2152 洛谷 2634 聪聪可可

题目

在一棵$n$个点的树上找两个可以相同的点，问这两点间路径和是3的倍数的概率

---

分析

树形dp，方程比较简单，但是求答案比较难，所以详见代码，时间复杂度$O(n)$

---

代码

#include <cstdio>
#include <cctype>
#define rr register
using namespace std;
const int N=20011;
struct node{int y,w,next;}e[N<<1];
int ls[N],dp[N][3],ans,n,k=1;
inline signed iut(){
	rr int ans=0; rr char c=getchar();
	while (!isdigit(c)) c=getchar();
	while (isdigit(c)) ans=(ans<<3)+(ans<<1)+(c^48),c=getchar();
	return ans;
}
inline signed gcd(int a,int b){return b?gcd(b,a%b):a;}
inline void dfs(int x,int fa){
	dp[x][0]=1;
	for (rr int i=ls[x];i;i=e[i].next)
	if (e[i].y!=fa){
		dfs(e[i].y,x);
		for (rr int j=0;j<3;++j)
		    ans+=dp[e[i].y][j]*dp[x][(6-e[i].w-j)%3]<<1;//意思就是该条路径上的答案乘不需要这条路径的答案，因为点对无序，所以要乘2
		for (rr int j=0;j<3;++j)
		    dp[x][(e[i].w+j)%3]+=dp[e[i].y][j];//就是传递答案
	}
}
signed main(){
	ans=n=iut();
	for (rr int i=1;i<n;++i){
		rr int x=iut(),y=iut(),w=iut()%3;
		e[++k]=(node){y,w,ls[x]},ls[x]=k,
		e[++k]=(node){x,w,ls[y]},ls[y]=k;
	}
	dfs(1,0);
	rr int all=n*n,Gcd=gcd(ans,all);
	return !printf("%d/%d",ans/Gcd,all/Gcd); 
}