题解 P2664 【树上游戏】
Salamander
2018-01-31 09:12:42
考虑点分治，对于当前分治中心，统计出它自己出发到分治块内的所有路径对自己答案的贡献，和经过它的路径对当前分治块内点的贡献。自己出发到分治块内的所有路径对自己答案的贡献很好求，现在考虑怎么求经过它的路径对当前分治块内点的贡献。
我们对于当前分治中心的每一个子树分别考虑，令$cnt[i]$为从分治中心出发的进入其他子树的所有路径中，包含颜色$i$的路径条数，$size$为除了该子树外当前分治块内所有的点的个数。那么，我们dfs这棵子树计算贡献，假设当前dfs到$x$，首先给$sum_x$加上$\sum cnt[i]$，即所有在其他子树中出现的颜色的贡献总和，然后计算$x$到分治中心的路径上颜色的贡献。
对于一个出现在分治中心到$x$的路径上的颜色$c$，它对$x$的贡献为$size-cnt[c]$，因为$c$已经在一些路径上出现，它现在能产生的额外贡献为它原来没有出现的路径条数。所以我们给$sum_x$还要加上$size-cnt[c]$。同时$c$也会对$x$的子树内所有点产生贡献，所以这个贡献要像标记一样往下传递，然后标记一下$c$的贡献已经被计算过，往下dfs时就不用再次计算了。
所以只需要对于当前分治中心求出$cnt$数组和每棵子树的$size$，进入一棵子树时减去子树自己内部对$cnt$产生的贡献。同时为了防止复杂度退化，我们不能对于所有颜色求$cnt$，要先统计一下当前分治块内有哪些颜色出现了，这样枚举块内所有颜色的复杂度才是O(分治块大小)。
复杂度$O(n\log n)$。
```cpp
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
#define For(i,_beg,_end) for(int i=(_beg),i##end=(_end);i<=i##end;++i)
#define Rep(i,_beg,_end) for(int i=(_beg),i##end=(_end);i>=i##end;--i)
template<typename T>T Max(const T &x,const T &y){return x<y?y:x;}
template<typename T>T Min(const T &x,const T &y){return x<y?x:y;}
template<typename T>int chkmax(T &x,const T &y){return x<y?(x=y,1):0;}
template<typename T>int chkmin(T &x,const T &y){return x>y?(x=y,1):0;}
template<typename T>void read(T &x){
T f=1;char ch=getchar();
for(;ch<'0'||ch>'9';ch=getchar())if(ch=='-')f=-1;
for(x=0;ch>='0'&&ch<='9';ch=getchar())x=x*10+ch-'0';
x*=f;
}
typedef long long LL;
const int maxn=100010,inf=0x7fffffff;
struct edge{
int to,nxt;
}e[maxn<<1];
int n,num,head[maxn],c[maxn];
int f[maxn],size[maxn],Size,root;
int has[maxn],col[maxn],top,cl[maxn],tp;
LL sum[maxn],res[maxn],cnt[maxn],ct[maxn],cct[maxn],path,tot;
bool vis[maxn],exist[maxn];
void addedge(int,int);
void get_sum(int,int);//计算子树大小
void get_root(int,int);//计算重心
void Dfs(int,int,LL*);//统计哪些颜色出现过，同时计算cnt数组
void Modify(int,int,LL);//更新子树的答案
void Calc(int);//对于分治中心统计答案
void Solve(int);//点分治
int main(){
read(n);
For(i,1,n) read(c[i]);
For(i,1,n-1){
int u,v;
read(u);read(v);
addedge(u,v);
}
Solve(1);
For(i,1,n) printf("%lld\n",sum[i]);
return 0;
}
void Calc(int x){
get_sum(x,Size=0);
tot=top=0;
Dfs(x,0,cnt);
For(i,1,top) exist[col[i]]=false;//要及时还原标记数组
For(i,1,tp=top){
tot+=cnt[cl[i]=col[i]];
cct[col[i]]=cnt[col[i]];
}
sum[x]+=tot;
LL temp=tot;
for(int i=head[x];i;i=e[i].nxt) if(!vis[e[i].to]){
has[c[x]]=1;top=0;
Dfs(e[i].to,x,ct);
has[c[x]]=0;
For(j,1,top) exist[col[j]]=false;
cnt[c[x]]-=size[e[i].to];
tot-=size[e[i].to];
For(j,1,top){
cnt[col[j]]-=ct[col[j]];
tot-=ct[col[j]];
}
path=size[x]-size[e[i].to];
Modify(e[i].to,x,0);
cnt[c[x]]+=size[e[i].to];
tot=temp;
For(j,1,top){
cnt[col[j]]=cct[col[j]];
ct[col[j]]=0;
}
}
For(i,1,tp) cnt[cl[i]]=0;//还原数组
vis[x]=true;
}
void Modify(int x,int fa,LL lst){//lst为之前下传下来的贡献和
LL tag=lst;
if(++has[c[x]]==1) tag+=path-cnt[c[x]];
sum[x]+=tot+tag;
for(int i=head[x];i;i=e[i].nxt){
if(e[i].to==fa||vis[e[i].to]) continue;
Modify(e[i].to,x,tag);
}
has[c[x]]--;
}
void Dfs(int x,int fa,LL *cnt){
if(!exist[c[x]]) col[++top]=c[x],exist[c[x]]=true;
if(++has[c[x]]==1) cnt[c[x]]+=size[x];
for(int i=head[x];i;i=e[i].nxt){
if(e[i].to==fa||vis[e[i].to]) continue;
Dfs(e[i].to,x,cnt);
}
has[c[x]]--;
}
void Solve(int x){
Calc(x);
for(int i=head[x];i;i=e[i].nxt) if(!vis[e[i].to]){
f[Size=root=0]=inf;
get_sum(e[i].to,x);
get_root(e[i].to,x);
Solve(root);
}
}
void get_sum(int x,int fa){
Size++;size[x]=1;
for(int i=head[x];i;i=e[i].nxt){
if(e[i].to==fa||vis[e[i].to]) continue;
get_sum(e[i].to,x);
size[x]+=size[e[i].to];
}
}
void get_root(int x,int fa){
f[x]=0;
for(int i=head[x];i;i=e[i].nxt){
if(e[i].to==fa||vis[e[i].to]) continue;
get_root(e[i].to,x);
chkmax(f[x],size[e[i].to]);
}
chkmax(f[x],Size-size[x]);
if(f[x]<f[root]) root=x;
}
void addedge(int u,int v){
e[++num].to=v;e[num].nxt=head[u];head[u]=num;
e[++num].to=u;e[num].nxt=head[v];head[v]=num;
}
```