hdu 6035 Colorful Tree(树形dp+技巧)

题目链接：hdu 6035 Colorful Tree

题意：

给你一棵树，每个节点有一种颜色，现在让你求所有点对的路径上不同的颜色数量的总和。

题解：

下面是官方题解：

单独考虑每一种颜色，答案就是对于每种颜色至少经过一次这种的路径条数之和。反过来思考只需要求有多少条路径没有经过这种颜色即可。直接做可以采用虚树的思想（不用真正建出来），对每种颜色的点按照 dfs 序列排个序，就能求出这些点把原来的树划分成的块的大小。这个过程实际上可以直接一次 dfs 求出。

这里的所说的单独考虑每种颜色，指的是对这个树开一个滤镜，比如对于这棵树，我只看颜色1，不是颜色1的点全部看为无色。

如果这棵树上的点全是颜色1，那么颜色1对答案的贡献就是这棵树的点对数n*(n-1)/2（因为每个点对的路径都有颜色1，就对的答案贡献了1）。但是并不是任意点对上的路径都有颜色1，所以对于不经过颜色1的联通块，我们要减去这个多算的联通块。

然后再推广到有n种颜色，每个颜色都要这样处理一下，这里可以用一个dfs直接一起处理全部颜色的情况，所以就能O(n)搞了。

#include<bits/stdc++.h>
#define F(i,a,b) for(int i=(a);i<=(b);++i)
using namespace std;
typedef long long ll;

const int N=2e5+7;
int n,c[N],x,y,cas,pos[N],cur[N];
vector<int>g[N];
ll ans,rem[N];

inline ll sum(ll x){return x*(x-1)>>1;}

int dfs(int x=1,int fa=0)
{
    int sz=1,szv,pre=pos[c[x]];
    pos[c[x]]=x;
    for(int &it:g[x])if(it!=fa)
    {
        cur[x]=0,szv=dfs(it,x);
        sz+=szv,ans-=sum(szv-cur[x]);
    }
    pre?cur[pre]+=sz:rem[c[x]]+=sz;
    pos[c[x]]=pre;
    return sz;
}

int main(){
    while(~scanf("%d",&n))
    {
        F(i,1,n)scanf("%d",c+i),g[i].clear(),rem[i]=0;
        F(i,2,n)
        {
            scanf("%d%d",&x,&y);
            g[x].push_back(y);
            g[y].push_back(x);
        }
        ans=sum(n)*n,dfs();
        F(i,1,n)ans-=sum(n-rem[i]);
        printf("Case #%d: %lld\n",++cas,ans);
    }
    return 0;
}