hdu 6035 树分治

题意：

一棵树，路径的权值等于路径上颜色的种类，问全部路径的价值和

 

思路：

题解的方法就不多讲了，这里提供一个树分治做法，首先要将原问题转换为每种颜色的经过路径数量的总和（也就是计算每种颜色的贡献，贡献为经过这种颜色的路径数量）

 

树分治的做法则为，维护重心下的子树第一次出现某种颜色的位置，通过这个位置已经这个点下面的子树大小，和一个相同颜色所有第一次出现的子树大小和来计算（注意去重），时间2s+通过，比较莽。。。毕竟比题解多个logn。。

 

代码：

 

#include<bits/stdc++.h>
using namespace std;
#define MEM(a,b) memset(a,b,sizeof(a))
#define PB push_back
typedef long long ll;
const int maxn =1e6+10;
ll n,k,ans;
int root,Max;
struct node{
    int v,next;
}edge[maxn*2];
int head[maxn],tot;
int si[maxn],maxv[maxn],vis[maxn];
int c[maxn],used[maxn],col[maxn],tmp[maxn];
ll sum,allson;
vector<int> us1,us2,su;
void init(){
    tot=ans=0;MEM(head,-1);MEM(vis,0);MEM(used,0);
}
void add_edge(int u,int v){
    edge[tot].v=v;edge[tot].next=head[u];head[u]=tot++;
}
void dfssi(int u,int f){
    si[u]=1;maxv[u]=0;
    for(int i=head[u];i!=-1;i=edge[i].next){
        int v=edge[i].v;
        if(v==f||vis[v])continue;
        dfssi(v,u);
        si[u]+=si[v];
        if(si[v]>maxv[u])maxv[u]=si[v];
    }
}
void dfsroot(int r,int u,int f){
    if(si[r]-si[u]>maxv[u]) maxv[u]=si[r]-si[u];
    if(maxv[u]<Max) Max=maxv[u],root=u;
    for(int i=head[u];i!=-1;i=edge[i].next){
        int v=edge[i].v;
        if(v==f||vis[v])continue;
        dfsroot(r,v,u);
    }
}
ll pre(int x,int fa){
    int ret=1;
    for(int i=head[x];i!=-1;i=edge[i].next){
        int v=edge[i].v;
        if(v!=fa&&!vis[v]) ret+=pre(v,x);
    }
    return ret;
}
ll dfsdis(int x, int fa) {
    ll ret=1;
    int cl=c[x];
    used[cl]++;
    for(int i=head[x];i!=-1;i=edge[i].next){
        int v=edge[i].v;
        if(v!=fa&&!vis[v]) ret+=dfsdis(v, x);
    }
    used[cl]--;
    if(used[cl]==0) ans+=ret*allson,tmp[cl]+=ret,us1.PB(cl),us2.PB(cl);
    return ret;
}
void cal(int x) {
    int id=0;
    us1.clear();su.clear();
    sum=allson=1;
    used[c[x]]++;
    for(int i=head[x];i!=-1;i=edge[i].next){
        int v=edge[i].v;
        if(vis[v]) continue;
        su.PB(pre(v,x));
        allson+=su[id++];
    }
    id=0;
    for(int i=head[x];i!=-1;i=edge[i].next){
        us2.clear();
        int v=edge[i].v;
        if(vis[v]) continue;
        allson-=su[id++];
        ll son=dfsdis(v, x);
        for(int u:us2)
            ans-=tmp[u]*col[u],col[u]+=tmp[u],tmp[u]=0;
        allson+=son;
        ans+=sum*son;
        sum+=son;
    }
    for(int u:us1) col[u]=0;
    used[c[x]]--;
}
void dfs(int u){
    Max=n;
    dfssi(u,0);
    dfsroot(u,u,0);
    vis[root]=1;
    cal(root);
    for(int i=head[root];i!=-1;i=edge[i].next){
        int v=edge[i].v;
        if(!vis[v])dfs(v);
    }
}

int main(){
    int ca=1;
    while(scanf("%d",&n)!=EOF){
        int u,v;
        init();
        for(int i=1;i<=n;i++) scanf("%d",&c[i]);
        for(int i=1;i<n;i++){
            scanf("%d%d",&u,&v);
            add_edge(u,v);add_edge(v,u);
        }
        dfs(1);
        printf("Case #%d: %lld\n",ca++,ans);
    }
    return 0;
}