bzoj #2152. 聰聰可可

題意

在一棵樹上任取兩點，問這兩點間路徑和爲3倍數的最簡概率

題解

樹形Dp，f[x][0/1/2]記錄x點以下（默認1爲根) mod 3 = 0/1/2 的邊的個數

\(\ f[x][0]\equiv0\)\(\pmod{3}\)

\(\ f[x][1]\equiv1\)\(\pmod{3}\)

\(\ f[x][2]\equiv2\)\(\pmod{3}\)

更新 ans += f[x][0] * f[cur][0] + f[x][1] * f[cur][2] + f[x][2] * f[cur][1]

f[x][0/1/2] += f[cur][0/1/2]

注意點

最後要求出真正ans，要約分

#include <cstdio>
#define maxn 20005

using namespace std;

int n, cnt, ans, head[maxn], f[maxn][3];
struct node{
	int next, to, v;
}e[maxn << 1];

void AddEdge(int u, int v, int dis){
	e[++cnt] = (node){head[u], v, dis}; head[u] = cnt;
	e[++cnt] = (node){head[v], u, dis}; head[v] = cnt;
}

int gcd(int x, int y){
	return (y == 0) ? x : gcd(y, x % y);
}

void dp(int x, int fa){
	f[x][0] = 1;
	for (int i = head[x]; i; i = e[i].next){
		int cur = e[i].to;
		if (cur != fa){
			dp(e[i].to, x);
			if (e[i].v % 3 == 1){
				int tmp = f[cur][0];
				f[cur][0] = f[cur][2];
				f[cur][2] = f[cur][1];
				f[cur][1] = tmp;
			}
			if (e[i].v % 3 == 2){
				int tmp = f[cur][0];
				f[cur][0] = f[cur][1];
				f[cur][1] = f[cur][2];
				f[cur][2] = tmp;
			}
			ans += f[x][0] * f[cur][0] + f[x][1] * f[cur][2] + f[x][2] * f[cur][1];
			f[x][0] += f[cur][0], f[x][1] += f[cur][1], f[x][2] += f[cur][2];
		}
	}
}

int main(){
	scanf("%d", &n);
	
	for (int x, y, z, i = 1; i < n; i++){
		scanf("%d%d%d", &x, &y, &z);
		AddEdge(x, y, z);
	}
	
	dp(1, 0);
	ans = ans * 2 + n; //(1, 2) -> (2, 1)    外加  (1, 1) ... (n, n) 
	int tmp = gcd(n * n, ans); // 通分 
	printf("%d/%d\n", ans / tmp, n * n / tmp);
	
	return 0;
}