正文
洛谷P2664 树上游戏（点分治）
小程序：扫一扫查出行【扫一扫了解最新限行尾号】
复制小程序
传送门
题解
因为一个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]$
不得不说思路真的挺妙的
//minamoto
#include<iostream>
#include<cstdio>
#include<cstring>
#define ll long long
using namespace std;
#define getc() (p1==p2&&(p2=(p1=buf)+fread(buf,1,1<<21,stdin),p1==p2)?EOF:*p1++)
char buf[<<],*p1=buf,*p2=buf;
template<class T>inline bool cmax(T&a,const T&b){return a<b?a=b,:;}
inline int read(){
#define num ch-'0'
char ch;bool flag=;int res;
while(!isdigit(ch=getc()))
(ch=='-')&&(flag=true);
for(res=num;isdigit(ch=getc());res=res*+num);
(flag)&&(res=-res);
#undef num
return res;
}
char sr[<<],z[];int C=-,Z;
inline void Ot(){fwrite(sr,,C+,stdout),C=-;}
inline void print(ll x){
if(C><<)Ot();if(x<)sr[++C]=,x=-x;
while(z[++Z]=x%+,x/=);
while(sr[++C]=z[Z],--Z);sr[++C]='\n';
}
const int N=;
int head[N],Next[N<<],ver[N<<],c[N],son[N],sz[N],size,cnt[N];
ll col[N],ans[N],much,sum,num;
int tot,n,rt;bool vis[N];
inline void add(int u,int v){
ver[++tot]=v,Next[tot]=head[u],head[u]=tot;
ver[++tot]=u,Next[tot]=head[v],head[v]=tot;
}
void findrt(int u,int fa){
sz[u]=,son[u]=;
for(int i=head[u];i;i=Next[i]){
int v=ver[i];
if(!vis[v]&&v!=fa){
findrt(v,u);
sz[u]+=sz[v];
cmax(son[u],sz[v]);
}
}
cmax(son[u],size-sz[u]);
if(son[u]<son[rt]) rt=u;
}
void dfs1(int u,int fa){
//要重新dfs一次，不能直接在找根时的树上做（不然子树和之类的会出错）
//顺便维护各种东西
sz[u]=,++cnt[c[u]];
for(int i=head[u];i;i=Next[i]){
int v=ver[i];
if(!vis[v]&&v!=fa){
dfs1(v,u),sz[u]+=sz[v];
}
}
if(cnt[c[u]]==) sum+=sz[u],col[c[u]]+=sz[u];
--cnt[c[u]];
}
void change(int u,int fa,int val){
++cnt[c[u]];
for(int i=head[u];i;i=Next[i]){
int v=ver[i];
if(!vis[v]&&v!=fa) change(v,u,val);
}
if(cnt[c[u]]==) sum+=sz[u]*val,col[c[u]]+=sz[u]*val;
--cnt[c[u]];
}
void dfs2(int u,int fa){
//把这棵子树里的颜色的影响消除掉
//顺便更新答案
++cnt[c[u]];
if(cnt[c[u]]==) sum-=col[c[u]],++num;
ans[u]+=sum+num*much;
for(int i=head[u];i;i=Next[i]){
int v=ver[i];
if(!vis[v]&&v!=fa) dfs2(v,u);
}
if(cnt[c[u]]==) sum+=col[c[u]],--num;
--cnt[c[u]];
}
void clear(int u,int fa){
cnt[c[u]]=col[c[u]]=;
for(int i=head[u];i;i=Next[i]){
int v=ver[i];
if(!vis[v]&&v!=fa) clear(v,u);
}
}
void did(int u){
//直接带进去乱搞
dfs1(u,);
ans[u]+=sum-col[c[u]]+sz[u];
for(int i=head[u];i;i=Next[i]){
int v=ver[i];
if(!vis[v]){
//dfs，然后把各种影响消除掉
++cnt[c[u]];
sum-=sz[v];
col[c[u]]-=sz[v];
change(v,u,-);
--cnt[c[u]];
much=sz[u]-sz[v];
dfs2(v,u);
++cnt[c[u]];
sum+=sz[v];
col[c[u]]+=sz[v];
change(v,u,);
--cnt[c[u]];
}
}
sum=,num=,clear(u,);
}
void solve(int u){
did(u),vis[u]=true;
for(int i=head[u];i;i=Next[i]){
int v=ver[i];
if(!vis[v]){
rt=,size=sz[v];
findrt(v,),solve(rt);
}
}
}
int main(){
n=read();
for(int i=;i<=n;++i) c[i]=read();
for(int i=;i<n;++i){
int u=read(),v=read();
add(u,v);
}
son[]=n+,rt=,size=n;
findrt(,),solve(rt);
for(int i=;i<=n;++i) print(ans[i]);
Ot();
return ;
}
