luogu P2664 树上游戏

本文介绍了一种使用点分治算法解决树形结构中节点颜色计数问题的方法。通过将树分解为多个子树,并利用桶记录颜色出现次数,实现了对所有节点到其它节点颜色数量的有效计算。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

背景:

h e h e . . . hehe... hehe...

题目传送门:

https://www.luogu.org/problemnew/show/P2664

题意:

一棵树,每一个点有一个颜色,定义 s i , j s_{i,j} si,j i i i j j j的颜色数量,求所有 s u m i = ∑ j = 1 n s i , j sum_i=\sum_{j=1}^{n}{s_{i,j}} sumi=j=1nsi,j

思路:

考虑点分治。
i i i j j j的颜色数量可以分为子树内的和子树外的贡献,用一个桶记录即可(因为一个颜色只有第一次才有贡献)。细节很多。

具体来说,对当前 R O O T ROOT ROOT的每个儿子的子树 d f s dfs dfs,如果当前点的颜色是 R O O T ROOT ROOT到当前点这条链上第一次出现,那么就把当前点的 s i z e size size加入桶。
先把所有儿子的子树全处理完,弄出来一个桶,注意根的颜色要特判。
然后统计答案,枚举根的儿子,先消除当前子树对桶的贡献,然后对当前子树 d f s dfs dfs,若当前点颜色第一次出现,就把当前颜色的桶的值改为 s i z e R O O T − s i z e x size_{ROOT}−size_x sizeROOTsizex x x x为当前儿子。
最后回溯时还原,每个子树统计完答案把影响加回来。

代码:

#include<cstdio>
#include<cstring>
#include<algorithm>
#define LL long long
using namespace std;
	int n,len=0;
	int col[200010],last[200010],size[200010],msize[200010];
	LL sum[200010],t[200010],cnt[200010];
	bool bz[200010];
	struct node{int x,y,next;} a[200010];
int SIZE,MIN,ROOT;
LL tot,del_size;
void ins(int x,int y)
{
	a[++len]=(node){x,y,last[x]}; last[x]=len;
}
void find_root(int x,int fa)
{
	size[x]=1;
	msize[x]=0;
	for(int i=last[x];i;i=a[i].next)
	{
		int y=a[i].y;
		if(y==fa||bz[y]) continue;
		find_root(y,x);
		size[x]+=size[y];
		msize[x]=max(msize[x],size[y]);
	}
	msize[x]=max(msize[x],SIZE-size[x]);
	if(MIN>msize[x]) MIN=msize[x],ROOT=x;
}
void solve(int x,int fa,int op)
{
	size[x]=1;
	cnt[col[x]]++;
	for(int i=last[x];i;i=a[i].next)
	{
		int y=a[i].y;
		if(y==fa||bz[y]) continue;
		solve(y,x,op),size[x]+=size[y];
	}
	cnt[col[x]]--;
	if(!cnt[col[x]]&&col[x]!=col[ROOT]) t[col[x]]+=size[x]*op,tot+=size[x]*op;
}
void clear(int x,int fa)
{
	t[col[x]]=cnt[col[x]]=0;
	for(int i=last[x];i;i=a[i].next)
	{
		int y=a[i].y;
		if(y==fa||bz[y]) continue;
		clear(y,x);
	}
}
void get_ans(int x,int fa)
{
	LL tmp=t[col[x]];
	if(!cnt[col[x]]&&col[x]!=col[ROOT]) t[col[x]]=del_size,tot=tot-tmp+del_size;
	sum[x]+=tot;
	cnt[col[x]]++;
	for(int i=last[x];i;i=a[i].next)
	{
		int y=a[i].y;
		if(y==fa||bz[y]) continue;
		get_ans(y,x);
	}
	cnt[col[x]]--;
	if(!cnt[col[x]]&&col[x]!=col[ROOT]) tot=tot-t[col[x]]+tmp,t[col[x]]=tmp;
}
void dfs(int x)
{
	bz[x]=true;
	clear(x,0);
	tot=0;
	for(int i=last[x];i;i=a[i].next)
	{
		int y=a[i].y;
		if(!bz[y]) solve(y,x,1);
	}
	del_size=1;
	for(int i=last[x];i;i=a[i].next) 
	{
		int y=a[i].y;
		if(!bz[y]) del_size+=size[y];
	}
	t[col[x]]=del_size;
	tot+=del_size;
	sum[x]+=tot;
	for(int i=last[x];i;i=a[i].next)
	{
		int y=a[i].y;
		if(bz[y]) continue;
		solve(a[i].y,x,-1);del_size-=size[y],t[col[x]]-=size[y],tot-=size[y];
		get_ans(y,x);
		solve(a[i].y,x,1);del_size+=size[y],t[col[x]]+=size[y],tot+=size[y];
	}
	for(int i=last[x];i;i=a[i].next)
	{
		int y=a[i].y;
		if(!bz[y]) SIZE=size[y],MIN=n,ROOT=0,find_root(y,x),dfs(ROOT);
	}
}
int main()
{
	int x,y;
	scanf("%d",&n);
	for(int i=1;i<=n;i++)
		scanf("%d",&col[i]);
	for(int i=1;i<n;i++)
	{
		scanf("%d %d",&x,&y);
		ins(x,y),ins(y,x);
	}
	SIZE=n,MIN=n,ROOT=0,find_root(1,0),dfs(ROOT);
	for(int i=1;i<=n;i++)
		printf("%lld\n",sum[i]);
}
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值