题解 P2664 【树上游戏】
y2823774827y
2019-02-25 14:58:52
刚学完点分治，"点分治真好玩"，$dalao$瞬间把这题丢过来，肛了三个小时终于$A$了
$$ans_i=\sum\limits_{j=1}^n sum(i,j)$$
其实就是一眼题系列，就是细节难处理
模板题告诉我们：要把其他子树都堆到一起进行计算，这题相对来说并不好处理
但其实是具有单调性的，设树的重心为$w$，一对父子关系的点对$(u,v)$，$v$是包含$u$的
更新：将重心看作一个无色节点，遍历其中一棵子树，能得到子树总贡献及每种颜色的贡献
统计答案：点对$(u,v)$，$v$统计在其他子树的贡献一定比$u$小，之前已得出其他子树每种颜色的贡献，减掉就好了
细节：
1. 统计重心贡献
1. 统计每种颜色贡献的标记记得弹出子树后清空
1. 实际菊花图能卡爆空间，可以用$map$处理数组，由于本题数据过水开子树序列$100$就能过了
```cpp
#include<bits/stdc++.h>
using namespace std;
typedef long long LL;
inline LL Read(){
LL x(0),f(1); char c=getchar();
while(c<'0' || c>'9'){ if(c=='-')f=-1; c=getchar(); }
while(c>='0' && c<='9') x=(x<<3)+(x<<1)+c-'0',c=getchar();
return x*f;
}
const LL maxn=1e5+9,inf=0x3f3f3f3f;
struct node{
LL to,nxt;
}dis[maxn<<1];
void Clear(LL u,LL now);
LL num,root,N,mi,tmp,n,tree,nsize;
LL head[maxn],size[maxn],a[maxn],cnt[100][maxn],ans[maxn],fir[maxn],V[maxn],sz[100],nsum[100],st[maxn];
bool visit[maxn];
inline void Add(LL u,LL v){
dis[++num]=(node){v,head[u]}, head[u]=num;
}
void Dfs(LL u){
size[u]=1; LL mx(0); visit[u]=true;
for(LL i=head[u];i;i=dis[i].nxt){
LL v(dis[i].to);
if(visit[v]) continue;
Dfs(v), size[u]+=size[v];
mx=max(mx,size[v]);
}mx=max(mx,N-size[u]);
if(mi>mx) mi=mx,root=u;
visit[u]=false;
}
void Up(LL u,LL now){
++sz[0], ++sz[now];
visit[u]=true;
bool f(false);
if(!fir[a[u]])
fir[a[u]]=f=true,
cnt[0][a[u]]+=size[u],cnt[now][a[u]]+=size[u],
nsum[0]+=size[u],nsum[now]+=size[u];
for(LL i=head[u];i;i=dis[i].nxt){
LL v(dis[i].to);
if(visit[v]) continue;
Up(v,now);
}visit[u]=false;
if(f) fir[a[u]]=false;
}
void Get(LL u,LL now){
bool f(false);
if(!fir[a[u]])
fir[a[u]]=f=true,tmp=tmp-cnt[0][a[u]]+cnt[now][a[u]],++tree;
ans[root]+=tree;
ans[u]+=tmp+tree*nsize; visit[u]=true;
for(LL i=head[u];i;i=dis[i].nxt){
LL v(dis[i].to);
if(visit[v]) continue;
Get(v,now);
}visit[u]=false;
if(f) --tree,fir[a[u]]=false,tmp=tmp+cnt[0][a[u]]-cnt[now][a[u]];
}
void Div(LL u){
mi=inf, Dfs(u);
Dfs(root);
visit[u=root]=true;
memset(nsum,0,sizeof(nsum));
nsize=1;
LL now(0);
for(LL i=head[u];i;i=dis[i].nxt){
LL v(dis[i].to);
if(visit[v]) continue;
++now;
Up(v,now);
}now=0;
for(LL i=head[u],S=nsize;i;i=dis[i].nxt){
LL v(dis[i].to);
if(visit[v]) continue;
++now;
nsize=sz[0]-sz[now]+1;
tmp=nsum[0]-nsum[now];
fir[a[u]]=true, tree=1;
tmp=tmp-cnt[0][a[u]]+cnt[now][a[u]];
Get(v,now);
fir[a[u]]=false;
}
now=0;
for(LL i=head[u];i;i=dis[i].nxt){
LL v(dis[i].to); if(visit[v]) continue;
++now;
Clear(v,now);
}
for(LL i=head[u],sum=N;i;i=dis[i].nxt){
LL v(dis[i].to);
if(visit[v]) continue;
N=size[v];
Div(v);
}
}
void Clear(LL u,LL now){
visit[u]=true;
cnt[0][a[u]]=cnt[now][a[u]]=0;
sz[0]=sz[now]=0;
nsum[0]=nsum[now]=0;
for(LL i=head[u];i;i=dis[i].nxt){
LL v(dis[i].to);
if(visit[v]) continue;
Clear(v,now);
}visit[u]=false;
}
int main(){
n=Read();
for(LL i=1;i<=n;++i) a[i]=Read();
for(LL i=1;i<n;++i){
LL u(Read()),v(Read());
Add(u,v), Add(v,u);
}
N=n, Div(1);
for(LL i=1;i<=n;++i) printf("%lld\n",ans[i]+1);
return 0;
}
```