[P2664]树上游戏(树上差分)

考虑对于每种颜色,如果把该颜色的节点全部删去的话,会形成很多小的树。对于一个仍然保留的节点,除了与它在同一颗小树上的点,其余点都会对它产生1的该颜色的贡献。那么一个节点的sum就是总颜色数*总节点数-Σ删掉每种颜色后该节点所属的树的大小。

于是考虑如何求新的树的大小,可以使用差分。对于一个颜色为c的节点u,它的儿子v的sum,便加上以v为根,最大的不含c颜色的树的大小。当在从v向它子树的某条路径上,找到了该路径上第一个颜色为c的点,那么把这个点的sum减去上面说的值。

具体如何实现,就对每种颜色记录一下当前不含该颜色的树的大小和这颗树的根节点。但不能每搜到一个节点都把与该节点颜色不同的树大小都更新一遍,所以我们对于每个颜色记录的是一个偏差值,另设一个计数器,偏差值和计数器的和为实际树的大小。当搜到节点u时,我们保存一下与u颜色相同的树的数值,然后计数器加一使其他颜色都加一,当前颜色清零。然后对于u的每个儿子,先设u的颜色的树根为该儿子,然后往下深搜,回溯时把该儿子的sum加上树的大小,然后树大小清零,继续算下一个儿子。

由于刚搜索到u时,我们可以知道它时那颗树的截止位置,但还没计算出该树的大小,所以不能用sum减,只记录一下该颜色树的根节点,都做完后再用sum减它的大小。注意原始树的根要单独处理一下。第二次更新sum时也要分类讨论要减的是不是原始树根,因为树根记录的是以它为根所有其他颜色的树的大小之和。

最后加一遍都好了。

#include
#include
#include
#define ll long long
using namespace std;
const int N=100010;
struct edge{
	int y,next;
}data[N*2];
struct color{
	int s,tot;
}col[N];
ll ans[N],ans1[N];
int n,m,c1[N],num,h[N],cnt[N],del[N],n1;
inline int read(){
	int x=0;char ch=getchar();
	while(ch<'0'||ch>'9')ch=getchar();
	while(ch>='0'&&ch<='9')x=x*10+ch-'0',ch=getchar();
	return x;
}
inline void addedge(int x,int y){
	data[++num].y=y;data[num].next=h[x],h[x]=num;
}
void dfs1(int u,int fa,int& tot){
	int tmp1=col[c1[u]].tot+tot,tmp2=col[c1[u]].s;
	++tot;col[c1[u]].tot=-tot;
	for(int i=h[u];i!=-1;i=data[i].next){
		int v=data[i].y;
		if(v!=fa){
			col[c1[u]].s=v;
			dfs1(v,u,tot);
			ans[v]=ans1[v]=1ll*(tot+col[c1[u]].tot),col[c1[u]].tot=-tot;
		}
	}
	del[u]=tmp2,col[c1[u]].tot=tmp1-tot,col[c1[u]].s=tmp2;
}
void getans(int u,int fa){
	for(int i=h[u];i!=-1;i=data[i].next){
		int v=data[i].y;
		if(v!=fa){ans[v]+=ans[u];getans(v,u);}
	}
}
int main(){
	n=read(),memset(cnt,0,sizeof cnt),m=0;
	for(int i=1;i<=n;++i)c1[i]=read(),++cnt[c1[i]];
	for(int i=1;i<=100000;++i)if(cnt[i]>0)++m;
	memset(h,-1,sizeof h),num=0;
	for(int x,y,i=1;i

 

你可能感兴趣的:(差分)