题意
给出一棵有n个结点的树，每个结点有一个颜色，问在这棵树的所有路径中，每条路径的颜色数求和是多少。
思路
求每种颜色的贡献可以转化为总的和减去每种颜色在哪些路径上没有出现的贡献，一个颜色在树块中的贡献为sz*(sz-1)/2。
这个题就是把一棵树分成许多个树块然后统计。
1号结点为一个所有颜色都有的点，特殊处理，否则当某种颜色在树的下面而且不是和根节点同一种颜色的时候，就无法进行计算了。
#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int INF = 0x3f3f3f3f;
const int N = 2e5 + 10;
vector<int> g[N];
int n, col[N], facol[N], fason[N];
LL num[N], root[N], sz[N];
/*
* 求每种颜色的贡献可以转化为总的和减去每种颜色在哪些路径上没有出现的贡献
* 一个颜色在树块中的贡献为sz*(sz-1)/2
* 这个题就是把一棵树分成许多个树块然后统计。
* 1号结点为一个所有颜色都有的点,这样可以方便统计
* facol[col[u]]表示和第u个点同样的颜色的离得最近的祖先的儿子是什么
* fason[u]是facol[col[u]]的简化,表示第u个点离得最近的祖先的儿子是什么
* num[u]表示u这个点的父结点的颜色在u这个分支块所统治的树块的大小
* root[col]表示第col种颜色与根节点的树块的大小
*/
void dfs(int u, int f) {
fason[u] = facol[col[u]];
sz[u] = 1;
for(int i = 0; i < g[u].size(); i++) {
int v = g[u][i]; if(v == f) continue;
facol[col[u]] = v;
dfs(v, u);
facol[col[u]] = fason[u];
sz[u] += sz[v];
// num[u] += sz[v];
}
// printf("%d : %d - %d - %lld - %lld - %lld
", u, col[u], fason[u], num[u], root[col[u]], sz[u]);
num[u] += sz[u];
if(fason[u] == 1) root[col[u]] -= sz[u];
else num[fason[u]] -= sz[u];
}
LL solve() {
for(int i = 0; i <= n; i++) g[i].clear();
for(int i = 1; i <= n; i++) scanf("%d", &col[i]);
for(int i = 1; i < n; i++) {
int u, v;
scanf("%d%d", &u, &v);
g[u].push_back(v);
g[v].push_back(u);
}
for(int i = 1; i <= n; i++) root[i] = n, facol[i] = 1, num[i] = 0;
dfs(1, 1);
LL ans = 1LL * n * n * (n - 1) / 2;
for(int i = 2; i <= n; i++) ans -= num[i] * (num[i] - 1) / 2;
for(int i = 1; i <= n; i++) ans -= root[i] * (root[i] - 1) / 2;
return ans;
}
int main() {
int cas = 1;
while(~scanf("%d", &n))
printf("Case #%d: %lld
", cas++, solve());
return 0;
}