5058. 【GDSOI2017模拟4.13】采蘑菇

Description

A君住在魔法森林里，魔法森林可以看做一棵n个结点的树，结点从1~n编号。树中的每个结点上都生长着蘑菇。蘑菇有许多不同的种类，但同一个结点上的蘑菇都是同一种类，更具体地，i号结点上生长着种类为c[i]的蘑菇。
现在A君打算出去采蘑菇，但他并不知道哪里的蘑菇更好，因此他选定起点s后会等概率随机选择树中的某个结点t作为终点，之后从s沿着(s,t)间的最短路径走到t.并且A君会采摘途中所经过的所有结点上的蘑菇。
现在A君想知道，对于每一个结点u，假如他从这个结点出发，他最后能采摘到的蘑菇种类数的期望是多少。为了方便，你告诉A君答案*n的值即可。

Input

第一行一个整数n表示结点数。
第二行n个整数c[i]表示每个结点的蘑菇的种类。
接下来n-1行每行两个数u[i],v[i]表示树中的一条边。

Output

输出n行每行一个整数，第i行的整数表示起点为结点i时的答案。

Sample Input

5
1 2 3 2 3
1 2
1 3
2 4
2 5

Sample Output

10
9
12
9
11

Data Constraint

30%的数据：n <= 2000
另有20%的数据：给出的第i条边为{i,i+1}
另有20%的数据：蘑菇的种类最多3种
100%的数据：1 <= n <= 3*10^5 , 0 <= c[i] <= n

Solutioin

这题可以用点分治或换根+线段树来做，这两种做法都是O(nlogn)的（听说还可以用虚树做）。但是，一位dalao却想出了一种巧妙的O(n)做法。

不妨把这棵树看成一个有根树。先定义几个量，令siz[x]表示以x为根的子树的大小，up[x]表示x到根的路径上离x最近的、父亲节点颜色为c[x]​的点的编号，sum[x]表示up值为x的点的siz总和，如下图所示（黑色的点表示颜色相同的点）：

考虑颜色c[x]​对那堆黄色点（黄色点的颜色不一定相同）的答案的贡献，发现从那堆黄色点出发，经过颜色c[x]​的路径有2种：第一种是向上走经过fa[ up[ x ] ]​​的路径，这样的贡献是n−siz[ up[ x ] ]​​；第二种是向下走经过点x的路径，这样的贡献是sum[ up[ x ] ]​​。因此，对于一个点x（这里的x即上图的up[ x ]），它对黄色点的两种贡献的总和为n−siz[ x ]+sum[ x ]​。因为up[ x ]对在x（这里的x指上图中的x）的子树里的点的答案没有贡献，所以在计算以x为根的子树的答案时，要减去n−siz[ up[ x ] ]+sum[ up[ x ] ]。

考虑特殊情况。当节点x没有up时，如下图所示：

显然，若一个点x到根的路径中不包含黑色节点，那么除去所有的黑色节点及其子树内的点，其他所有的点都可以获得以   (没有up的黑色节点 ) 为根的子树的大小的和的贡献。

此外，还没有考虑起点的颜色贡献，因此答案要加上n。

实现的时候建议用dfs序把树转化成一个序列，然后用差分处理。

Code

#include<cstdio> 
#include<cstring>
#include<algorithm>
#include<cmath>
#define I int
#define ll long long
#define son t[k]
#define F(i,a,b) for(I i=a;i<=b;i++)
#define Fd(i,a,b) for(I i=a;i>=b;i--)
#define N 300005
using namespace std;
I n,a[N],x,y,dfn[N],sz[N],up[N],b[N],cnt;
I t[N<<1],nx[N<<1],ls[N],nex[N],las[N],tot;
ll sum[N],f[N];
void add(I x,I y){t[++tot]=y,nx[tot]=ls[x],ls[x]=tot;}
void ins(I x,I y,ll z){f[x]+=z,f[y+1]-=z;}
void dfs(I x,I y){
	I now=b[a[x]];
	up[x]=now;dfn[x]=++cnt;sz[x]=1;
	for(I k=ls[x];k;k=nx[k]) if(son!=y){
		dfs(b[a[x]]=son,x);
		sz[x]+=sz[son];
	}
	b[a[x]]=now;sum[up[x]]+=sz[x];
}
void dg(I x,I y){
	ins(dfn[x],dfn[x]+sz[x]-1,n-sz[x]+sum[x]);
	if(up[x]) ins(dfn[x],dfn[x]+sz[x]-1,-n+sz[up[x]]-sum[up[x]]);
	else{nex[x]=las[a[x]],las[a[x]]=x;}
	for(I k=ls[x];k;k=nx[k]) if(son!=y) dg(son,x);
}
I main(){
	freopen("mushroom.in","r",stdin);
	freopen("mushroom.out","w",stdout);
	scanf("%d",&n);
	F(i,1,n) scanf("%d",&a[i]);
	add(0,1),add(1,0);
	F(i,1,n-1){
		scanf("%d%d",&x,&y);
		add(x,y),add(y,x);
	}
	dfs(1,0);
	dg(1,0);
	F(i,0,n) if(las[i]){
		sum[0]=0;
		for(I k=las[i];k;k=nex[k]) sum[0]+=sz[k];
		ins(1,n,sum[0]);
		for(I k=las[i];k;k=nex[k]) ins(dfn[k],dfn[k]+sz[k]-1,-sum[0]);
	}
	F(i,1,n) f[i]+=f[i-1];
	F(i,1,n) printf("%lld\n",f[dfn[i]]+n);
	return 0;
}