题目大意
给定一颗树,然后每个结点有一种颜色,然后求所有路径中不同颜色的个数之和。
思路
可以这样想,先假设所有路径中都包含所有的n种颜色,则可以得出ans = n * n * (n - 1) / 2
然后,减去路径中不存在某个颜色的路径条数,可以想象一下一棵树中某个联通分量中没有某个颜色,然后这个联通分量点的个数为cnt那么就需要在ans中减cnt * (cnt - 1) / 2
代码
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn = 2e5 + 10;
int n, col[maxn], sz[maxn], sum[maxn];
ll ans;
vector<int> G[maxn];
void dfs(int u, int p) {
sz[u] = 1;
int all = 0;
for(int i = 0; i < G[u].size(); i++) {
int v = G[u][i];
if(v == p) continue;
int pre = sum[col[u]];
dfs(v, u);
sz[u] += sz[v];
int temp = sz[v] - (sum[col[u]] - pre);
all += temp;
ans -= 1LL * temp * (temp - 1) / 2;
}
sum[col[u]]+=all+1;
}
bool vis[maxn];
int main() {
int cas = 0;
while(~scanf("%d", &n)) {
for(int i = 1; i <= n; i++) {
scanf("%d", &col[i]);
G[i].clear();
sum[i] = 0;
}
for(int i = 0; i < n - 1; i++) {
int u,v;
scanf("%d%d", &u, &v);
G[u].push_back(v);
G[v].push_back(u);
}
ans = 1LL * n * n * (n - 1) / 2;
dfs(1, -1);
for(int i = 1; i <= n; i++) {
ans -= 1LL * (n - sum[i]) * (n - sum[i] - 1) / 2;
}
printf("Case #%d: %lld\n",++cas, ans);
}
return 0;
}