树上点对问题很容易就想到点分治
点分治已经用掉一个log了 那么我们的问题就转化为：
在 o(n) 时间内求出 以该根为lca的所有的点对的答案！
1、对于树上的一个i 如果i的颜色在i到根的链上是第一次出现 那么可以对树上其他的点j （不考虑特殊情况）产生贡献siz[i] （j与子树i的每一个点都产生贡献1）
2、dfs求出sum（整棵树所有的贡献 和color[i] 这个颜色所有的贡献）
3、显然 根的答案为 ans[root]+=sum 也就是一端在根上的链
4、遍历子树 统计答案
但是有一些不是树上最短路径的贡献不能被统计进去 所以要在该子树统计答案的时候 删除该子树的影响 然后再回溯回去
5、设Ｘ＝sigma color[j 到根上（不包括根）的所有颜色] 由于这些颜色已经出现过，我们不能在该子树外计算其贡献）
设num为ｊ到根上（不包括根）的颜色数
设Ｙ为size[root]-size[该子树]（即所有其他子树＋根的点数）
则ans[j]+＝sum-X+num*Y;
这样一道黑题就大功告成啦
#include<bits/stdc++.h> using namespace std; #define rep(i,a,b) for(int i=(a);i<=(b);i++) #define repp(i,a,b) for(int i=(a);i>=(b);--i) #define ll long long #define see(x) (cerr<<(#x)<<'='<<(x)<<endl) #define inf 0x3f3f3f3f #define CLR(A,v) memset(A,v,sizeof A) ////////////////////////////////// const int N=2e6+10; int head[N],pos,n,m,v[N],siz[N],maxson[N],root,color[N],cnt[N],siz_tree,vis[N],a,b; ll sum,ans[N],much,num; struct Edge{int to,nex;}edge[N<<1]; void add(int a,int b){edge[++pos]=(Edge){b,head[a]};head[a]=pos;} void getroot(int x,int fa) { siz[x]=1;maxson[x]=0; for(int i=head[x];i;i=edge[i].nex) { int v=edge[i].to; if(vis[v]||v==fa)continue; getroot(v,x);siz[x]+=siz[v]; maxson[x]=max(maxson[x],siz[v]); } maxson[x]=max(maxson[x],siz_tree-siz[x]); if(maxson[x]<maxson[root])root=x; } void dfs1(int x,int fa) { cnt[v[x]]++;siz[x]=1; for(int i=head[x];i;i=edge[i].nex) { int v=edge[i].to; if(v==fa||vis[v])continue; dfs1(v,x);siz[x]+=siz[v]; } if(cnt[v[x]]==1) { sum+=siz[x]; color[v[x]]+=siz[x]; } cnt[v[x]]--; } void dfs2(int x,int fa) { cnt[v[x]]++; if(cnt[v[x]]==1)sum-=color[v[x]],num++; ans[x]+=sum+num*much; for(int i=head[x];i;i=edge[i].nex) { int v=edge[i].to; if(vis[v]||v==fa)continue; dfs2(v,x); } if(cnt[v[x]]==1)sum+=color[v[x]],num--; cnt[v[x]]--; } void change(int x,int fa,int value) { cnt[v[x]]++; for(int i=head[x];i;i=edge[i].nex) { int v=edge[i].to; if(v==fa||vis[v])continue; change(v,x,value); } if(cnt[v[x]]==1) { sum+=1ll*value*siz[x]; color[v[x]]+=1ll*value*siz[x]; } cnt[v[x]]--; } void clear(int x,int fa) { cnt[v[x]]=color[v[x]]=0; for(int i=head[x];i;i=edge[i].nex) { int v=edge[i].to; if(v==fa||vis[v])continue; clear(v,x); } } void calc(int x) { dfs1(x,0); ans[x]+=sum; for(int i=head[x];i;i=edge[i].nex) { int to=edge[i].to; if(vis[to])continue; cnt[v[x]]++; sum-=siz[to]; color[v[x]]-=siz[to]; change(to,x,-1); much=siz[x]-siz[to]; cnt[v[x]]--; dfs2(to,x); cnt[v[x]]++; sum+=siz[to]; color[v[x]]+=siz[to]; change(to,x,1);//回溯 cnt[v[x]]--; } sum=num=0; clear(x,0); } void solve(int x) { if(vis[x])return ;vis[x]=1; calc(x); for(int i=head[x];i;i=edge[i].nex) { int v=edge[i].to; if(vis[v])continue; root=0; siz_tree=siz[v]; getroot(v,x); solve(root); } } int main() { scanf("%d",&n); rep(i,1,n)scanf("%d",&v[i]); rep(i,1,n-1)scanf("%d%d",&a,&b),add(a,b),add(b,a); siz_tree=maxson[0]=n; root=0; getroot(1,0); solve(root); rep(i,1,n)printf("%lld ",ans[i]); return 0; }