【洛谷P2664】 树上游戏 点分治

code:

#include <bits/stdc++.h>   
 
#define N 200009   
 
#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],A[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;  
 
}        
  
int ou;     
 
ll tmp,tot,bu[N];           
 
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;      
 
    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 u,int ff) 
{          
    ll tt=bu[val[u]];           
 
    tmp=tmp-bu[val[u]]+ou;            
 
    bu[val[u]]=ou;                    

    Sum[u]+=tmp;       

    for(int i=hd[u];i;i=nex[i]) 
    {
        int v=to[i]; 
 
        if(vis[v]||v==ff)    continue;          
 
        calc_v(v,u);                       
    }  

    tmp=tmp-bu[val[u]]+tt;    

    bu[val[u]]=tt;    
}
void clr(int u,int ff)
{
    cn[u].clear();    
    bu[val[u]]=0;         
    for(int i=hd[u];i;i=nex[i])
    {
        int v=to[i]; 
        if(v==ff||vis[v])   continue;   
        clr(v,u);     
    }
}
void calc(int u) 
{         
    tot=0;       
 
    getnode(u,u,0,0);       

    for(int i=hd[u];i;i=nex[i])           
    {
        int v=to[i];   
 
        if(vis[v])    continue;     
 
        // memset(cnt,0,sizeof(cnt));  
 
        get_col(v,v,u);    
 
        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=siz[u]-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; 
 
        bu[val[u]]=ou;  
 
        calc_v(v,u);     
  
        bu[val[u]]=tt;   
 
        for(it=cn[v].begin();it!=cn[v].end();it++) bu[it->first]+=it->second;      
  
    }          
 
    clr(u,0);     
}
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]), A[i]=val[i]; 
 
    sort(A+1,A+1+n); 
 
    for(i=1;i<=n;++i)     val[i]=lower_bound(A+1,A+1+n,val[i])-A;     
 
    for(i=1;i<n;++i) 
    {
        int u,v; 
 
        scanf("%d%d",&u,&v),add(u,v),add(v,u);  
    }          
 
    sn=mx[0]=n,root=0,getroot(1,0),dfs(root);   
 
    for(i=1;i<=n;++i)     printf("%lld\n",Sum[i]);        
 
    return 0; 
}