2017 Multi-University Training Contest 1003(hdu 6035) Colorful Tree

看了很多人的博客才弄懂这一道题目，感觉他们解释的还不够详细。

题意：

一棵有n个点的树，树上每个点都有颜色c[i]，保证每两个点之间的路径只有一条，定义每条路径的值为这条路径上经过的不同颜色数量和。求所有路径的值。

题解：

主要思想就是对于每个颜色单独考虑。

先把所有的情况都算上，也就是对于n个点，一共有n*（n-1）条不同的路径。那么对于n种颜色，一共就是ans = n*n*（n-1）种情况。

然后呢，对于当前颜色 c[i]，如果有一个节点数为 gsize[dx] 子图不包含颜色c[i]，那就把ans -= gsize[dx]*(gsize[dx]-1)/2，减去所有不包含颜色c[i]的子图的种类，剩下的就是正确答案。

下图中就是对于颜色2，一共有4个子图需要排除。



然后我们dfs的时候是从上往下，一次得到答案，具体的实现在代码中。

代码：

#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <stack>
#include <bitset>
#include <queue>
#include <set>
#include <map>
#include <string>
#include <algorithm>
using namespace std;
#define clr(a,b) memset(a,b,sizeof(a))
#define pb(a) push_back(a)
#define fir first
#define se second
#define LL long long
typedef pair<int,int> pii;
typedef pair<LL,int> pli;
typedef pair<LL,LL> pll;
const int inf = 0x3f3f3f3f;
const int maxn = 2e5+5;
LL mod = 1e9+7;
double eps = 0.00000001;
double PI = acos(-1);

int c[maxn];  //节点的颜色
LL sum[maxn],ans,gsize[maxn];
//sum[i] 颜色i的节点数 , gsize[i] 以节点i为根的节点数
vector<int> g[maxn]; //图

void dfs(int u,int fa) {
LL pre = sum[c[u]];
gsize[u] = 1;
LL add = 1;
for(auto v : g[u]) {
if(v == fa) continue;
dfs(v,u);
gsize[u] += gsize[v];
LL tmp = gsize[v] - (sum[c[u]] - pre);
ans -= tmp*(tmp-1)/2;
pre = sum[c[u]];
add += tmp;
}
sum[c[u]] += add;
}

int main() {
LL n;int ca = 1;
while(scanf("%I64d",&n) != EOF) {
for(int i = 1;i <= n;i++) {  //初始化
sum[i] = gsize[i] = 0;
g[i].clear();
}
for(int i = 1;i <= n;i++) scanf("%d",&c[i]);
for(int i = 1;i < n;i++) {
int u,v;
scanf("%d%d",&u,&v);
g[u].pb(v);
g[v].pb(u);
}
ans = n*n*(n-1)/2;  //假设每种颜色存在于每一条路径中
dfs(1,0);
for(int i = 1;i <= n;i++) {
LL tmp = n-sum[i];
ans -= tmp*(tmp-1)/2;
}
printf("Case #%d: %I64d\n",ca++,ans);
}
}