题解
因为一个sb错误调了一个晚上……鬼晓得我为什么$solve(rt)$会写成$solve(v)$啊!!!一个$O(logn)$被我硬生生写成$O(n)$了竟然还能过$5$个点……话说还一直以为只有动态点分会很难没想到一般点分都这么可啪……%%%大佬
我们考虑一下,对于一棵树,我们要处理的是子树对根的答案的贡献,以及经过根的路径的贡献(也就是$LCA$为根的点对的答案)。
对于树中的一个点$i$,如果$i$的颜色是在$i$到根的路径上第一次出现,那么所有与$i$的$LCA$为根的点,都能与$i$的子树形成点对,且$i$的颜色会对那些点产生产生贡献$size[i]$(先不考虑其他子树中是否有相同颜色)。
那么我们考虑$dfs$一遍整棵树,并记录,所有颜色产生的贡献$col[i]$以及贡献总和$sum$。然后$dfs$子树,并把其中的所有颜色的贡献给减掉。然后考虑子树中的一个点,设$x$为到根上的所有点的贡献总和,$num$表示根到节点上的颜色个数,$y=size[root]-size[该子树]$,那么$ans[j]+=sum-x+num*size[y]$
然后最后记得加上对根节点的答案的贡献$ans[root]+=sum-col[根的颜色]+size[root]$
不得不说思路真的挺妙的
1 //minamoto 2 #include<iostream> 3 #include<cstdio> 4 #include<cstring> 5 #define ll long long 6 using namespace std; 7 #define getc() (p1==p2&&(p2=(p1=buf)+fread(buf,1,1<<21,stdin),p1==p2)?EOF:*p1++) 8 char buf[1<<21],*p1=buf,*p2=buf; 9 template<class T>inline bool cmax(T&a,const T&b){return a<b?a=b,1:0;} 10 inline int read(){ 11 #define num ch-'0' 12 char ch;bool flag=0;int res; 13 while(!isdigit(ch=getc())) 14 (ch=='-')&&(flag=true); 15 for(res=num;isdigit(ch=getc());res=res*10+num); 16 (flag)&&(res=-res); 17 #undef num 18 return res; 19 } 20 char sr[1<<21],z[20];int C=-1,Z; 21 inline void Ot(){fwrite(sr,1,C+1,stdout),C=-1;} 22 inline void print(ll x){ 23 if(C>1<<20)Ot();if(x<0)sr[++C]=45,x=-x; 24 while(z[++Z]=x%10+48,x/=10); 25 while(sr[++C]=z[Z],--Z);sr[++C]=' '; 26 } 27 const int N=200005; 28 int head[N],Next[N<<1],ver[N<<1],c[N],son[N],sz[N],size,cnt[N]; 29 ll col[N],ans[N],much,sum,num; 30 int tot,n,rt;bool vis[N]; 31 inline void add(int u,int v){ 32 ver[++tot]=v,Next[tot]=head[u],head[u]=tot; 33 ver[++tot]=u,Next[tot]=head[v],head[v]=tot; 34 } 35 void findrt(int u,int fa){ 36 sz[u]=1,son[u]=0; 37 for(int i=head[u];i;i=Next[i]){ 38 int v=ver[i]; 39 if(!vis[v]&&v!=fa){ 40 findrt(v,u); 41 sz[u]+=sz[v]; 42 cmax(son[u],sz[v]); 43 } 44 } 45 cmax(son[u],size-sz[u]); 46 if(son[u]<son[rt]) rt=u; 47 } 48 void dfs1(int u,int fa){ 49 //要重新dfs一次,不能直接在找根时的树上做(不然子树和之类的会出错) 50 //顺便维护各种东西 51 sz[u]=1,++cnt[c[u]]; 52 for(int i=head[u];i;i=Next[i]){ 53 int v=ver[i]; 54 if(!vis[v]&&v!=fa){ 55 dfs1(v,u),sz[u]+=sz[v]; 56 } 57 } 58 if(cnt[c[u]]==1) sum+=sz[u],col[c[u]]+=sz[u]; 59 --cnt[c[u]]; 60 } 61 void change(int u,int fa,int val){ 62 ++cnt[c[u]]; 63 for(int i=head[u];i;i=Next[i]){ 64 int v=ver[i]; 65 if(!vis[v]&&v!=fa) change(v,u,val); 66 } 67 if(cnt[c[u]]==1) sum+=sz[u]*val,col[c[u]]+=sz[u]*val; 68 --cnt[c[u]]; 69 } 70 void dfs2(int u,int fa){ 71 //把这棵子树里的颜色的影响消除掉 72 //顺便更新答案 73 ++cnt[c[u]]; 74 if(cnt[c[u]]==1) sum-=col[c[u]],++num; 75 ans[u]+=sum+num*much; 76 for(int i=head[u];i;i=Next[i]){ 77 int v=ver[i]; 78 if(!vis[v]&&v!=fa) dfs2(v,u); 79 } 80 if(cnt[c[u]]==1) sum+=col[c[u]],--num; 81 --cnt[c[u]]; 82 } 83 void clear(int u,int fa){ 84 cnt[c[u]]=col[c[u]]=0; 85 for(int i=head[u];i;i=Next[i]){ 86 int v=ver[i]; 87 if(!vis[v]&&v!=fa) clear(v,u); 88 } 89 } 90 void did(int u){ 91 //直接带进去乱搞 92 dfs1(u,0); 93 ans[u]+=sum-col[c[u]]+sz[u]; 94 for(int i=head[u];i;i=Next[i]){ 95 int v=ver[i]; 96 if(!vis[v]){ 97 //dfs,然后把各种影响消除掉 98 ++cnt[c[u]]; 99 sum-=sz[v]; 100 col[c[u]]-=sz[v]; 101 change(v,u,-1); 102 --cnt[c[u]]; 103 much=sz[u]-sz[v]; 104 dfs2(v,u); 105 ++cnt[c[u]]; 106 sum+=sz[v]; 107 col[c[u]]+=sz[v]; 108 change(v,u,1); 109 --cnt[c[u]]; 110 } 111 } 112 sum=0,num=0,clear(u,0); 113 } 114 void solve(int u){ 115 did(u),vis[u]=true; 116 for(int i=head[u];i;i=Next[i]){ 117 int v=ver[i]; 118 if(!vis[v]){ 119 rt=0,size=sz[v]; 120 findrt(v,0),solve(rt); 121 } 122 } 123 } 124 int main(){ 125 n=read(); 126 for(int i=1;i<=n;++i) c[i]=read(); 127 for(int i=1;i<n;++i){ 128 int u=read(),v=read(); 129 add(u,v); 130 } 131 son[0]=n+1,rt=0,size=n; 132 findrt(1,0),solve(rt); 133 for(int i=1;i<=n;++i) print(ans[i]); 134 Ot(); 135 return 0; 136 }