写到一半 ~

code: 

#include <bits/stdc++.h>   
#define N 1009   
#define ll long long 
#define setIO(s) freopen(s".in","r",stdin)   
using namespace std;  
ll Sum[N];
int n,edges,root,sn;     
int val[N],hd[N],to[N<<1],nex[N<<1],size[N],mx[N],vis[N];  
inline void add(int u,int v) 
{
    nex[++edges]=hd[u],hd[u]=edges,to[edges]=v;  
}     
void getroot(int u,int ff) 
{
    size[u]=1,mx[u]=0; 
    for(int i=hd[u];i;i=nex[i]) 
    {
        int v=to[i];   
        if(v==ff||vis[v])    continue;    
        getroot(v,u);  
        size[u]+=size[v];   
        mx[u]=max(mx[u],size[v]);  
    }
    mx[u]=max(mx[u],sn-size[u]); 
    if(mx[u]<mx[root])   root=u;  
}        

ll tmp,tot,bu[N];         
int tl,ou;       
map<int,ll>cn[N];     
map<int,ll>::iterator it;       
int dep[N],cnt[N],siz[N];                
void getnode(int top,int u,int ff,int cur) 
{                    
    if(!cnt[val[u]])   
        ++cur;      
    ++cnt[val[u]];  
    Sum[top]+=(ll)cur;          
    siz[u]=1;   
    dep[u]=dep[ff]+1;     
    for(int i=hd[u];i;i=nex[i])  
    {
        int v=to[i]; 
        if(v==ff||vis[v])     continue;    
        getnode(top,v,u,cur);  
        siz[u]+=siz[v];  
    }      
    --cnt[val[u]];                            
}            
void get_col(int top,int u,int ff) 
{                 
    if(!cnt[val[u]])     cn[top][val[u]]+=(ll)siz[u];             
    ++cnt[val[u]];     
    for(int i=hd[u];i;i=nex[i]) 
    {
        int v=to[i]; 
        if(v==ff||vis[v])   continue;      
        get_col(top,v,u); 
    } 
    --cnt[val[u]];        
}          
void calc_v(int top,int u,int ff) 
{          
    ll tt=bu[val[u]];           
    tmp=tmp-bu[val[u]]+ou+dep[u];                
    bu[val[u]]=ou+dep[u]; 
    Sum[u]+=tmp;       
    for(int i=hd[u];i;i=nex[i]) 
    {
        int v=to[i]; 
        if(vis[v]||v==ff)    continue;                        
    }     
    bu[val[u]]=tt;    
}
void calc(int u) 
{         
    tot=0;  
    int totsize=0;        
    getnode(u,u,0,0);                   // 处理完 u 点答案     
    for(int i=hd[u];i;i=nex[i])           
    {
        int v=to[i];   
        if(vis[v])    continue;        
        get_col(v,v,u);   
        totsize+=siz[v];   
        for(it=cn[v].begin();it!=cn[v].end();it++)    
        {
            tot+=it->second; 
            bu[it->first]+=it->second;                
        }
    }                                 
    for(int i=hd[u];i;i=nex[i])    
    {   
        int v=to[i];   
        if(vis[v])    continue;       
        tmp=tot;     
        ou=totsize-siz[v];   
        for(it=cn[v].begin();it!=cn[v].end();it++)    
        {
            bu[it->first]-=it->second;  
            tmp-=it->second;            
        }
        ll tt=bu[val[u]];  
        tmp=tmp-bu[val[u]]+ou+1;                
        bu[val[u]]=ou+1;                    
        calc_v(v,v,u);    
        bu[val[u]]=tt;        
    }         
    memset(bu,0,sizeof(bu));  
    memset(cnt,0,sizeof(cnt));    
    for(int i=0;i<N;++i)  cn[i].clear();    
}
void dfs(int u) 
{
    calc(u);      
    vis[u]=1;  
    for(int i=hd[u];i;i=nex[i]) 
    {
        int v=to[i];   
        if(vis[v])    continue;     
        root=0,sn=size[v],getroot(v,u),dfs(root);   
    }
}
int main() 
{ 
    // setIO("input");  
    int i,j;        
    scanf("%d",&n);     
    for(i=1;i<=n;++i)     scanf("%d",&val[i]);     
    for(i=1;i<n;++i) 
    {
        int u,v; 
        scanf("%d%d",&u,&v),add(u,v),add(v,u);  
    }          
    mx[0]=N,root=0,getroot(1,0);        

    return 0; 
}