【Luogu2664】树上游戏（点分治）

【Luogu2664】树上游戏（点分治）

题面

洛谷

题解

很好的一道点分治题。
首先直接点分治，考虑过每个分治重心的链的贡献。
我们从分治重心开始找每种颜色，强制令一种颜色只在其到分治重心的链上第一次出现的位置统计贡献，假设子树大小是(size)，那么对于当前分治重心的其他所有子树都会产生(size)的贡献。
那么考虑当前分治重心每个子树的每个点会得到的贡献，首先把这棵子树内的贡献删去，然后记录其他所有颜色的贡献和。如果当前颜色在这棵子树内第一次出现，那么其他所有子树都必定会产生贡献，那么把这部分贡献加上，继续递归就行了。
详细的看看代码吧。

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<cmath>
#include<algorithm>
#include<vector>
using namespace std;
#define ll long long
#define pb push_back
#define MAX 100100
inline int read()
{
	int x=0;bool t=false;char ch=getchar();
	while((ch<'0'||ch>'9')&&ch!='-')ch=getchar();
	if(ch=='-')t=true,ch=getchar();
	while(ch<='9'&&ch>='0')x=x*10+ch-48,ch=getchar();
	return t?-x:x;
}
vector<int> E[MAX];
int n,c[MAX];ll ans[MAX];
int mx,Size,rt,sz[MAX];
bool vis[MAX];
void Getroot(int u,int ff)
{
	sz[u]=1;int ret=0;
	for(int v:E[u])
	{
		if(v==ff||vis[v])continue;
		Getroot(v,u);sz[u]+=sz[v];
		ret=max(ret,sz[v]);
	}
	ret=max(ret,Size-sz[u]);
	if(ret<mx)mx=ret,rt=u;
}
int cnt[MAX];ll num[MAX],sum;
void dfs(int u,int ff,int opt)
{
	if(!cnt[c[u]]++)num[c[u]]+=opt*sz[u],sum+=opt*sz[u];
	for(int v:E[u])if(v!=ff&&!vis[v])dfs(v,u,opt);
	--cnt[c[u]];
}
void dfs(int u,int ff)
{
	if(!cnt[c[u]]++)sum+=Size-num[c[u]];
	ans[u]+=sum;
	for(int v:E[u])if(v!=ff&&!vis[v])dfs(v,u);
	if(!--cnt[c[u]])sum-=Size-num[c[u]];
}
void Divide(int u)
{
	vis[u]=true;Getroot(u,0);
	dfs(u,0,1);ans[u]+=sum;Size=sz[u];
	for(int v:E[u])
	{
		if(vis[v])continue;
        num[c[u]]-=sz[v];sum-=sz[v];Size-=sz[v];
        cnt[c[u]]=1;dfs(v,u,-1);cnt[c[u]]=0;
        dfs(v,u);
        cnt[c[u]]=1;dfs(v,u,1);cnt[c[u]]=0;
        num[c[u]]+=sz[v];sum+=sz[v];Size+=sz[v];
	}
	dfs(u,0,-1);
	for(int v:E[u])
	{
		if(vis[v])continue;
		mx=Size=sz[v];Getroot(v,u);
		Divide(rt);
	}
}
int main()
{
	n=read();
	for(int i=1;i<=n;++i)c[i]=read();
	for(int i=1,u,v;i<n;++i)u=read(),v=read(),E[u].pb(v),E[v].pb(u);
	mx=Size=n;Getroot(1,0);
	Divide(rt);
	for(int i=1;i<=n;++i)printf("%lld
",ans[i]);
	return 0;