HDU 6015 Colorful Tree（2017多校）
题目传送门：http://acm.hdu.edu.cn/showproblem.php?pid=6035
题意：
给出一颗树， 树上路径权值等于路径上不同颜色的数量，求所有路径权值之和；
解题思路：
大佬讲的很清楚了：https://blog.csdn.net/qust1508060414/article/details/76360368
AC code：
1 #include <bits/stdc++.h> 2 #define INF 0x3f3f3f3f 3 #define LL long long 4 using namespace std; 5 const int MAXN = 2e5+10; 6 int c[MAXN]; 7 bool vis[MAXN]; 8 vector<int>e[MAXN]; 9 LL sum[MAXN], Size[MAXN]; 10 LL ans; 11 12 void dfs(int x, int fa) 13 { 14 Size[x] = 1; 15 sum[c[x]]++; 16 LL pre = sum[c[x]]; 17 int len = e[x].size(); 18 for(int i = 0; i < len; i++){ 19 if(e[x][i] == fa) continue; 20 dfs(e[x][i], x); 21 Size[x]+=Size[e[x][i]]; 22 LL Count = Size[e[x][i]]-(sum[c[x]]-pre); 23 ans = ans+(Count*(Count-1))/2; 24 sum[c[x]]+=Count; 25 pre = sum[c[x]]; 26 } 27 } 28 29 int main() 30 { 31 int N, cas = 1; 32 LL num = 0, ct = 0; 33 while(~scanf("%d", &N)){ 34 num = 0; ans = 0; 35 memset(sum, 0, sizeof(sum)); 36 memset(vis, 0, sizeof(vis)); 37 for(int i = 1; i <= N; i++){ 38 e[i].clear(); 39 scanf("%d", &c[i]); 40 if(!vis[c[i]]){ 41 vis[c[i]] = true; 42 num++; 43 } 44 } 45 int u, v; 46 for(int i = 1; i < N; i++){ 47 scanf("%d %d", &u, &v); 48 e[u].push_back(v); 49 e[v].push_back(u); 50 } 51 dfs(1, -1); 52 LL res = (LL)(1LL*N*(N-1))/2*num; 53 for(int i = 1; i <= N; i++){ 54 if(vis[i]){ 55 ct = N-sum[i]; 56 ans+=ct*(ct-1)/2; 57 } 58 } 59 printf("Case #%d: %lld\n", cas++, res-ans); 60 } 61 return 0; 62 }
