JZOJ5058. 采蘑菇（O(n)算法，慎入）

JZOJ5058. 【GDSOI2017模拟4.13】采蘑菇

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

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

这道题有虚数做法，有换根线段树做法，也有点分治做法。
它们（除了虚树，因为我并不会）都带有一只log。虽然这道题时限比较宽松，但是有一位dalao想出了O(n)的做法。

我们随便取一个点为根。
定义一下：
$s{z}_x$为以$x$为根的子树的大小，$u{p}_x$为从$x$到根的路径上，离$x$最近的且颜色与$x$相同的祖先的与$x$在同一条路径上的儿子。

假设我们做到节点$u$，颜色为$c$。考虑经过$x$的路径的贡献。
黑色的节点表示与u同色的节点。
$f$为离$x$最近的且颜色与$x$相同的祖先，$x$为$u{p}_u$。

假设以$x$为根的子树内的节点为起点，其它节点为终点，那么很显然以$x$为根的子树内每个节点都能接受到其它节点个数的贡献，也就是$n-s{z}_x$。同时另$su{m}_x$表示$up$值为$x$的节点，以它们为根的子树大小的总和，贡献加上$su{m}_x$。

但是不难发现，这样会算重。
以$u$为例，它的子树内，路径上有颜色$c$的贡献应该至少来自$u$。同理，它也没有$su{m}_x$的贡献。相当于这一次它们什么贡献都没有获得。

总结的来说，$x$的子树内，我们先全部加上$n-s{z}_x+su{m}_x$。同时对于$u$及和它等概念的节点（例如它旁边那个小黑），它们并没有贡献。但是为了方便，不操作等价于再减去$n-s{z}_{u{p}_u}+su{m}_{u{p}_u}$。

那么还有一种情况，节点$u$没有$up$。

我们直接统计这类节点的，以它们为根的子树大小，总和记为$V$。
那么不属于这些节点的点可以获得$V$的贡献，属于这些节点的点没有获得贡献。

最后，因为没有记以自己为起点的贡献，所以都加上$n$。

对于所有的操作，在$dfs$序上，一次操作对应的点都是连续的，所以可以用差分来实现。时间复杂度$O(N)$。

最后再次膜拜蔡$BirdBird$，想出这个神仙做法。
并十分感谢Z_Y_S为我这个蒟蒻讲解这个神仙算法。

#include
#include
#include
using namespace std;
const int N=3e5+1000;
int n;
int c[N+1];
vector<int> e[N+1],e1[N+1],s[N+1];
int siz[N+1],id[N+1],up[N+1],sum[N+1];
void dfs(int u,int fa)
{
	siz[u]=1,id[u]=++id[0];
	up[u]=s[c[u]][s[c[u]].size()-1];
	if(up[u]==0)
		e1[c[u]].push_back(u);
	for(int i=0;i<e[u].size();i++)
	{
		int v=e[u][i];
		if(v==fa)
			continue;
		s[c[u]].push_back(v);
		dfs(v,u);
		s[c[u]].pop_back();
		siz[u]+=siz[v];
	}
	sum[up[u]]+=siz[u];
}
int tree[N+1];
void update(int a,int b,int c){tree[a]+=c,tree[b+1]-=c;}
void solve(int u,int fa)
{
	update(id[u],id[u]+siz[u]-1,n-siz[u]+sum[u]);
	if(up[u])
		update(id[u],id[u]+siz[u]-1,-n+siz[up[u]]-sum[up[u]]);
	for(int i=0;i<e[u].size();i++)
	{
		int v=e[u][i];
		if(v==fa)
			continue;
		solve(v,u);
	}
}
int num[N+1],ans[N+1];
bool cmp(int x,int y)
{
	return c[x]<c[y];
}
int main()
{
	freopen("mushroom.in","r",stdin);
	freopen("mushroom.out","w",stdout);
	scanf("%d",&n);
	for(int i=1;i<=n;i++)
		scanf("%d",&c[i]),s[i].push_back(0);
	s[0].push_back(0);
	for(int i=1;i<n;i++)
	{
		int a,b;
		scanf("%d %d",&a,&b);
		e[a].push_back(b),e[b].push_back(a);
	}
	dfs(1,0),solve(1,0);
	for(int i=0;i<=n;i++)
	{
		int v=0;
		for(int j=0;j<e1[i].size();j++)
			v+=siz[e1[i][j]];
		update(1,n,v);
		for(int j=0;j<e1[i].size();j++)
			update(id[e1[i][j]],id[e1[i][j]]+siz[e1[i][j]]-1,-v);
	}
	for(int i=1;i<=n;i++)
		ans[i]=ans[i-1]+tree[i];
	for(int i=1;i<=n;i++)
		printf("%d\n",ans[id[i]]+n);
	return 0;
}

Z_Y_Sdalao的代码QWQ。