首页 > itarticle > colorful tree

colorful tree

admin 11月 14, 2020 0

题目链接

题意

一棵有n个节点的树，每个点有一个点权，问在这棵树上所有的$frac{n(n-1)}{2}$条链颜色种类的和是多少。

思路

对于每种颜色计算其贡献，我们可以先算出每种颜色在多少条链中不出现，用总数减一下就可以得到在多少条链中出现，dfs序处理一下，遍历到一个点的时候我只考虑这个点与其儿子构成的连同块，在这些块里面都不出现这个颜色。

dfs

void (int u, int f)
{
	sz[u] = 1;
	ll all = 0;
	for (auto &v : G[u])
	{
		if (v == f)
			continue;
		ll last = sum[c[u]];
		dfs(v, u);
		sz[u] += sz[v];
		ll add = sum[c[u]] - last;
		ans += (sz[v] - add) * (sz[v] - add - 1) / 2;
		all += sz[v] - add;
	}
	sum[c[u]] += all + 1;
}

int main()
{
	for (int i = 1; i <= n; i++)
	{
		if (!hasc[i])
			continue;
		ans += (n - sum[i]) * (n - sum[i] - 1) / 2ll;
	}
	printf("%lldn", (n - 1) * n / 2 * cnt - ans);
	return 0;
}