【luoguP2664】树上游戏

题目链接

题解

考虑对于一个节点，其所有的路径分为两种：经过根节点和不经过根节点。运用点分治就可以每次计算对于当前树而言，经过根的总贡献。而剩余没有经过根的部分，则在删到根节点形成新的树之后递归求解，这样一定是可以做到不重不漏的。
接下来，就该考虑如何统计贡献。假设先将根节点定为起点，那么在一条路径中，某种颜色出现的位置深度最小的节点会对根节点产生其子树大小的贡献，这是因为从它往后的每个节点都已包含这个颜色，且后续不需要重复统计。那么，我们先将根节点得到的贡献计算出，然后再想办法更新它的子节点。
在这里我采用的方法是，用(cnt_i)记录(dfs)过程中颜色(i)出现的次数，在回溯时相应地修改回去。若访问到某个节点时其颜色的(cnt)为0，说明它是这条路径上这个颜色第一次出现的位置，计入贡献。
好了，这样我们用一次(dfs)将根节点的贡献更新好了，那么子节点的贡献怎么更新呢？
首先牢记我们每次只更新经过根节点的路径的贡献，这里的根节点是指点分治过程中一次更新所指认的根。考虑子节点到根节点的路径再加上其余的子树构成了一棵新的树。设当前子节点为(u)，它的深度为1的祖先为(fa)，当前子节点的颜色为(c_u)，原本的根节点为(root)，(weight_i)表示不在(fa)的子树中的颜色(i)的贡献。(c_u)有这么几种可能：

1. 没有出现过，对于当前节点的贡献为(size_{root}-size_{fa}+1)；
2. 出现过，且在到根节点的路径上没有出现，对于当前点的贡献要加上(size_{root}-size_{fa}+1-weight_{c_u})；
3. 出现过，且最晚是在到根节点的路径上出现的，对于当前点的贡献不变且为(size_{root}-size_{fa}+1)。

于是(u)得到的贡献即为其直属父亲的贡献加上它的颜色带来的贡献差值。
代码中有注释。

代码

[click]

#include <cstdio>
#include <cctype>
#include <cstdlib>
#include <ctime>
typedef long long ll;
const int maxn=1e5+10;
int head[maxn],to[maxn<<1],nxt[maxn<<1];
int c[maxn],val[maxn],size[maxn],cnt[maxn],weight[maxn];
int tot,root;
ll s[maxn];
bool vis[maxn];

int max(int x,int y) {return x>y?x:y;}
int read()
{
	int res=0;
	char ch=getchar();
	while(!isdigit(ch))
		ch=getchar();
	while(isdigit(ch))
		res=res*10+ch-'0',ch=getchar();
	return res;
}
void add(int u,int v)
{
	nxt[++tot]=head[u];
	head[u]=tot;
	to[tot]=v;
}
void build(int u,int fa,int n)
{
	int maxSize=0;
	size[u]=1;
	for (int i=head[u];i;i=nxt[i])
	{
		int v=to[i];
		if (vis[v]||v==fa)
			continue;
		build(v, u, n);
		size[u]+=size[v];
		maxSize=max(maxSize, size[v]);
	}
	maxSize=max(maxSize, n-size[u]);
	if (!root&&maxSize<=n/2)
		root=u;
}
void dfs(int u,int fa)
{
	val[u]=0;
	if (!cnt[c[u]])//说明这是在到根的路径上该颜色出现深度最小的位置 
		val[u]+=size[u],weight[c[u]]+=size[u];
	cnt[c[u]]++;
	for (int i=head[u];i;i=nxt[i])
	{
		int v=to[i];
		if (vis[v]||v==fa)
			continue;
		dfs(v, u);
		val[u]+=val[v];
	}
	cnt[c[u]]--;
}
void deal(int u,int fa,int p)
{
	if (!cnt[c[u]])
		weight[c[u]]+=p*size[u];
	cnt[c[u]]++;
	for (int i=head[u];i;i=nxt[i])
	{
		int v=to[i];
		if (vis[v]||v==fa)
			continue;
		deal(v, u, p);
	}
	cnt[c[u]]--;
}
void update(int u,int fa,int sum,ll d,int tmp)//每到一个节点就视为整棵树是由该点到根的路径+其它子树所构成的树，且当前节点为根 
{
	int pre=weight[c[u]];
	d+=tmp-pre;
	weight[c[u]]=tmp;
	//weight[i]记为起点为当前点，终点不在当前子树的包含颜色i的路径
	//其最大值为除掉该节点所在的root的子树的节点数之和，即tmp 
	s[u]+=d;
	for (int i=head[u];i;i=nxt[i])
	{
		int v=to[i];
		if (vis[v]||v==fa)
			continue;
		update(v, u, sum+1, d, tmp);
	}
	weight[c[u]]=pre;
}
void solve(int u)
{
	dfs(u, u);//val[v]表示v的子树对u产生的贡献
	cnt[c[u]]++;
	for (int i=head[u];i;i=nxt[i])
	{
		int v=to[i];
		if (vis[v])
			continue;
		deal(v, u, -1),weight[c[u]]-=size[v];
		update(v, u, size[u]-size[v]+1, val[u]-val[v]-size[v], size[u]-size[v]);
		deal(v, u, 1),weight[c[u]]+=size[v];
	} 
	cnt[c[u]]--;
	deal(u, 0, -1);
	s[u]+=val[u];
}
void work(int u)
{
	vis[u]=1;
	solve(u);
	for (int i=head[u];i;i=nxt[i])
	{
		int v=to[i];
		if (vis[v])
			continue;
		root=0;
		build(v, 0, size[v]);
		build(root, 0, size[v]);
		work(root);
	}
}
int main()
{
	int n=read();
	for (int i=1;i<=n;i++)
		c[i]=read();
	for (int i=1;i<=n-1;i++)
	{
		int u=read(),v=read();
		add(u, v),add(v, u);
	}
	srand((unsigned)time(0));
	build(rand()%n+1, 0, n);
	build(root, 0, n);
	work(root);
	for (int i=1;i<=n;i++)
		printf("%lld
",s[i]);
	return 0;
}