gym/102253C Colorful Tree 树上计数

题意：有一颗数，每个点有一个颜色，定义两点之间的距离为两点路径之间不同颜色的数目，问所有路径的距离和是多少？

思路：每个颜色的贡献为路径中有这个颜色的路径数。先假设所有路径都会经过一种颜色，再减去不会经过这个颜色的路径数就是这个颜色的贡献。

看一下这个博客会好理解一些吧：https://blog.csdn.net/f_zyj/article/details/76168180

代码：

#include <bits/stdc++.h>
#define LL long long
using namespace std;
const int maxn = 200010;
vector<int> G[maxn];
int c[maxn];
LL ans[maxn * 2], res;
int s[maxn];
int n;
void add(int x, int y) {
	G[x].push_back(y);
	G[y].push_back(x);
}
LL dfs(int x, int fa) {
	LL sz = 1, tmp = s[c[x]];
	s[c[x]] = x;
	for (int i = 0; i < G[x].size(); i++) {
		int y = G[x][i];
		if(y == fa) continue;
		ans[x] = 0;
		LL sz1 = dfs(y, x);
		res -= (LL)(sz1 - ans[x]) * (sz1 - ans[x] - 1) / 2;
		sz += sz1;
	}
	if(tmp <= n) ans[tmp] += sz;
	else ans[tmp] -= sz;
	s[c[x]] = tmp;
	return sz;
}
set<int> st;
set<int>::iterator it;
int main() {
	int x, y, kase = 0;
	while(~scanf("%d", &n)) {
		st.clear();
		res = 0;
		for (int i = 1; i <= n; i++) {
			scanf("%d", &c[i]);
			G[i].clear();
			s[i] = 0;
			st.insert(c[i]);
			ans[i] = 0;
		}
		for (it = st.begin(); it != st.end(); it++) {
			ans[(*it) + n] = n;
			s[(*it)] = (n + (*it));
		}
		for (int i = 1; i < n; i++) {
			scanf("%d%d", &x, &y);
			add(x, y);
		}
		res = (LL) st.size() * n * (n - 1) / 2;
		dfs(1, -1);
		for (it = st.begin(); it != st.end(); it++) {
			res -= ans[(*it) + n] * (ans[(*it) + n] - 1) / 2;
		}
		printf("Case #%d: %lld\n", ++kase, res);
	}
}