P2664 树上游戏

我太蒻了，这题好难……

首先进行点分治，分裂时自己分裂的子树对自己的节点可以利用深搜搜出来，如果一个颜色在此节点到根的链上第一次出现，那么对根的贡献为其子树大小，如果不是第一次出现，那么无贡献。

然后考虑子树对子树的贡献，那么先减去自己子树做的贡献，然后搜索整棵子树，对于自己到根上的所有点的贡献都赋为整棵树的大小减去这棵子树的大小，然后求解即可，最后记得回溯自己子树的贡献。

然后是点分治最重要的东东，清空数组。

对于颜色计数 \(cnt_i\)，在回溯时会自动清空，总贡献是一个变量，直接情况即可，子树大小不需要清空，因为每次搜索时会赋初值 \(1\)，最后是数组 \(c_i\)，代表颜色 \(i\) 做出的贡献，感觉不是很好清空，那么再搜索一遍整棵子树，将节点颜色对应的 \(c_i\) 清空。

复杂度：\(O(n\times\log(n))\)，废话。

#include<iostream>
#include<vector>
#include<cstdio>
#define int long long
using namespace std;
const int N=1e5+5;
int n,root,mp[N],siz[N],tot,vis[N],ans[N],c[N],cnt[N],sum,nbr,t[N],lim;
vector<int>a[N];
void findzx(int x,int fa)
{
	mp[x]=0;
	siz[x]=1;
	int len=a[x].size();
	for(int i=0;i<len;i++)
	{
		if(a[x][i]==fa||vis[a[x][i]])continue;
		findzx(a[x][i],x);
		mp[x]=max(mp[x],siz[a[x][i]]);
		siz[x]+=siz[a[x][i]];
	}
	mp[x]=max(mp[x],tot-siz[x]);
	if(mp[x]<mp[root])root=x;
}
void dfs(int x,int fa)
{
	cnt[t[x]]++;
	siz[x]=1;
	int len=a[x].size();
	for(int i=0;i<len;i++)
	{
		if(vis[a[x][i]]||a[x][i]==fa)continue;
		dfs(a[x][i],x);
		siz[x]+=siz[a[x][i]];
	}
	if(cnt[t[x]]==1)
	{
		c[t[x]]+=siz[x];
		sum+=siz[x];
	}
	cnt[t[x]]--;
}
void change(int x,int fa,int val)
{
	cnt[t[x]]++;
	int len=a[x].size();
	for(int i=0;i<len;i++)
	{
		if(vis[a[x][i]]||a[x][i]==fa)continue;
		change(a[x][i],x,val);
	}
	if(cnt[t[x]]==1)
	{
		c[t[x]]+=val*siz[x];
		sum+=val*siz[x];
	}
	cnt[t[x]]--;
}
void dfs2(int x,int fa)
{
	cnt[t[x]]++;
	if(cnt[t[x]]==1)
	{ 
		sum-=c[t[x]];
		lim++;
	}
	ans[x]+=lim*nbr+sum;
	int len=a[x].size();
	for(int i=0;i<len;i++)
	{
		if(a[x][i]==fa||vis[a[x][i]])continue;
		dfs2(a[x][i],x);
	}
	if(cnt[t[x]]==1)
	{
		sum+=c[t[x]];
		lim--;
	}
	cnt[t[x]]--;
}
void clea(int x,int fa)
{
	cnt[t[x]]=c[t[x]]=0;
	int len=a[x].size();
	for(int i=0;i<len;i++)
	{
		if(vis[a[x][i]]||a[x][i]==fa)continue;
		clea(a[x][i],x);
	}
}
void clac(int x)
{
	dfs(x,0);
	ans[x]+=sum-c[t[x]]+siz[x];
	int len=a[x].size();
	for(int i=0;i<len;i++)
	{
		if(vis[a[x][i]])continue;
		cnt[t[x]]++;
		
		c[t[x]]-=siz[a[x][i]];
		sum-=siz[a[x][i]];
		change(a[x][i],x,-1);
		cnt[t[x]]--;
		nbr=siz[x]-siz[a[x][i]];
		dfs2(a[x][i],x);
		cnt[t[x]]++;
		change(a[x][i],x,1);
		c[t[x]]+=siz[a[x][i]];
		sum+=siz[a[x][i]];
		cnt[t[x]]--;
	}
	sum=lim=0;
	clea(x,0);
}
void divide(int x,int num)
{
	clac(x);
	vis[x]=1;
	int len=a[x].size();
	for(int i=0;i<len;i++)
	{
		if(vis[a[x][i]])continue;
		if(siz[x]>siz[a[x][i]])tot=siz[a[x][i]];
		else tot=num-siz[x];
		root=0;
		findzx(a[x][i],x);
		divide(root,tot);
	}
}
signed main()
{
	int n;
	scanf("%lld",&n); 
	for(int i=1;i<=n;i++)scanf("%lld",&t[i]);
	for(int i=1;i<n;i++)
	{
		int u,v;
		scanf("%lld%lld",&u,&v);
		a[u].push_back(v);
		a[v].push_back(u);
	}
	mp[0]=1e9;
	tot=n;
	findzx(1,0);
	divide(root,n);
	for(int i=1;i<=n;i++)printf("%lld\n",ans[i]);
	return 0;
}