hdu 6035 Colorful Tree(dfs)

题意：

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

解题思路：

可以把问题转化为对每种颜色有多少条不同的路径至少经过这种颜色的点，然后加和。

求有多少条路径经过可以转换为总路径数-没有经过的路径数，只要求出没有经过的路径数就好了。

对于每种颜色没有经过自己的路径条数我们可以通过一遍dfs求得。一个点u往下dfs的时候，对于自己的一个儿子v的子树，找出其中和u颜色相同距离自己最近的点i，把这些点的size都减去，省下的就是一个和u颜色都不同的点的连通块了，而这些点组成的路径是不会穿过u对应的颜色的，就求出来了。

具体的过程需要通过一个sum数组，sum[i]数组保存的其实就是在自己到根节点这条路径上高度最高的颜色为i点的size的和。每次访问一颗子树前先把当前的size[i]保存在pre,跑完子树后比较一下sum[i]和pre，sum[i]-pre的值就是子树中距离u最近的颜色相同的点的size和，减去这些点后，就求出连通块的元素个数了。

#include <bits/stdc++.h>
#define ps push_back
using namespace std;

const int maxn=2e5+5;
vector<int>edg[maxn];
int siz[maxn];
int book[maxn];
int c[maxn];
int sum[maxn];
int tot;
int pos[maxn];
int dfn[maxn];
long long ans;

void dfs(int x, int fa)
{
siz[x]=1;
int i, v, t=c[x], pre=sum[c[x]];

int add=0;
for(i=0; i<(int)edg[x].size(); i++)
{
v=edg[x][i];
if(v==fa)continue;
dfs(v, x);
siz[x]+=siz[v];
//在以v为根节点的子树中离v最近的与v颜色相同的子树节点之和
long long  temp=siz[v]-(sum[c[x]]-pre);
pre=sum[c[x]];
add+=temp;
ans-=temp*(temp-1LL)/2;
}
sum[c[x]]+=add+1;
}

int main()
{
int i, j, e=1;
long long n;
while(~scanf("%lld", &n))
{
memset(sum, 0, sizeof sum);
for(i=1; i<=n; i++)
{
scanf("%d", &c[i]);
siz[i]=0;
}
int x, y;
for(i=2; i<=n; i++)
{
scanf("%d%d", &x, &y);
edg[x].ps(y);
edg[y].ps(x);
}

ans=(long long)n*n*(n-1)/2;
dfs(1, 0);
for(i=1; i<=n; i++)
{
long long temp=n-sum[i];
ans-=temp*(temp-1)/2;
edg[i].clear();
}
printf("Case #%d: %lld\n", e++, ans);

}
}