【YBTOJ】【Luogu P2634】[国家集训队]聪聪可可

链接:

洛谷

题目大意：

在一棵 \(n\) 个节点的树上随机选两个点（可以选同一个），求两点路径长度和是 \(3\) 的倍数的期望个数。

正文:

点分治经典题。在板子的基础上，求两点之间路径长度时可以不断模三，那么对答案的统计时，\(ans=2t_1\cdot t_2+{t_0}^2\)，其中 \(t_i\) 表示路径长度和模三为 \(i\) 的点对数量。

代码：

const int N = 100010;

int n;

struct edge
{
	int to, nxt, val;
}e[N << 1];
int head[N], tot;

void add(int u, int v, int w)
{
	e[++tot] = (edge) {v, head[u], w}, head[u] = tot;
}

bool vis[N];
int rt, sum, ans;
int sz[N], son[N];

void getRoot (int u, int fa)
{
	sz[u] = 1; son[u] = 0;
	for (int i = head[u]; i; i = e[i].nxt)
	{
		int v = e[i].to;
		if (vis[v] || v == fa) continue;
		getRoot(v, u);
		sz[u] += sz[v];
		son[u] = max(son[u], sz[v]);
	}
	son[u] = max(son[u], sum - sz[u]);
	if (son[u] < son[rt]) rt = u;
}

int dis[N], dissort[4];

void getDis (int u, int fa)
{
	dissort[dis[u]] ++;
	for (int i = head[u]; i; i = e[i].nxt)
	{
		int v = e[i].to;
		if (vis[v] || v == fa) continue;
		dis[v] = (dis[u] + e[i].val) % 3;
		getDis(v, u);
	}
}

int calc (int u, int val)
{
	dis[u] = val;
	dissort[0] = dissort[1] = dissort[2] = 0;
	getDis(u, 0);
	return dissort[0] * dissort[0] + 2 * dissort[1] * dissort[2];
}

void solve (int u)
{
	ans += calc(u, 0);
	vis[u] = 1;
	for (int i = head[u]; i; i = e[i].nxt)
	{
		int v = e[i].to;
		if (vis[v]) continue;
		ans -= calc (v, e[i].val);
		
		// Figure 1:
		sum = sz[v], son[0] = n, rt = 0;
		getRoot (v, u), solve(rt);
	}
}

int gcd(int a, int b)
{
	return b? gcd(b, a % b): a;
}

int main()
{
	scanf ("%d", &n);
	for (int i = 1; i < n; i++)
	{
		int u, v, w;
		scanf ("%d%d%d", &u, &v, &w); w %= 3;
		add(u, v, w), add(v, u, w);
	}
	son[0] = sum = n, getRoot(1, 0), solve(rt);
	int g = gcd(ans, n * n);
	printf ("%d/%d\n", ans / g, n * n / g);
	return 0;
}