HDU6035 Colorful Tree

题目连接：http://acm.hdu.edu.cn/showproblem.php?pid=6035

【题意】给定一棵有n个节点的树，ci表示节点i的颜色。任意两个节点间存在一条路径，该路径的权值为路径上不同颜色节点的数量，求解所有路径权值的和。

【分析】比赛期间想要通过树分治对答案进行求解，无奈发现始终无法完成对路径具体数值的快速统计和避免同色记重。赛后看了题解才发现自己还是太弱了。分析可以发现每个颜色对答案的贡献是经过该颜色的路径数。求解经过一个颜色的路径数还是不好求解，所以再转个方向去求解不经过该颜色的路径数。求解不经过对应颜色的路径数等于删除所有该颜色的节点后求解剩下的所有的联通块路径数的和。简单的删点重构树显然会T，所以采用虚树的概念进行求解，即不对树进行重构。此时需要利用dfs序，从而快速获取子树对应的区间，而区间大小就是节点数，可以快速求解路径数。具体方法看代码。代码中通过In保存对应节点dfs先序遍历值，Out保存对应先序遍历值的dfs序后续遍历值，合起来对应了以该节点为根的子树区间。

【代码】

#include<iostream>
#include<cstring>
#include<cstdio>
#include<algorithm>
#include<cmath>
#include<vector>
#include<fstream>
using namespace std;
#define Uint unsigned int
#define LL long long
vector<int> color[200200];
int In[200200];
int Out[200200];
int c[200200];
vector<int> edge[200200];
int cnt;
void dfs(int u,int f){
    In[u]=++cnt;
    for(int i=0;i<edge[u].size();++i){
        int v=edge[u][i];
        if(v==f)
            continue;
        dfs(v,u);
    }
    Out[In[u]]=cnt;
}
bool cmp(int x,int y){
    return In[x]<In[y];
}
LL Count(int idx,int s,int t,int l,int r){
    LL tot=r-l+1;
    LL ans=0;
    while(s<=t){
        int head=In[color[idx][s]]+1,tail=Out[In[color[idx][s]]];
        tot-=tail-head+2;
        s++;
        while(head<=tail){
            int he=head;
            int ta=Out[he];
            int i=s;
            while(i<=t && ta>=In[color[idx][i]])
                i++;
            ans+=Count(idx,s,i-1,he,ta);
            s=i;
            head=ta+1;
        }
    }
    ans+=tot*(tot-1)/2;
    return ans;
}

int main(){
    int n,u,v,cas=1;
    while(~scanf("%d",&n)){
        for(int i=1;i<=n;++i){
            In[i]=Out[i]=0;
            color[i].clear();
            edge[i].clear();
        }
        for(int i=1;i<=n;++i){
            scanf("%d",&c[i]);
            color[c[i]].push_back(i);
        }
        for(int i=1;i<n;++i){
            scanf("%d %d",&u,&v);
            edge[u].push_back(v);
            edge[v].push_back(u);
        }
        cnt=0;
        dfs(1,0);
        LL ans=0;
        for(int i=1;i<=n;++i){
            sort(color[i].begin(),color[i].end(),cmp);
            ans+=n*1ll*(n-1ll)/2-Count(i,0,color[i].size()-1,1,n);
        }
        printf("Case #%d: %I64d\n",cas++,ans);
    }
}