hdu6035 Colorful Tree 点分治

http://acm.hdu.edu.cn/showproblem.php?pid=6035

对着题面思考了1小时，看题解研究了1小时代码才看懂，菜哭.jpg

这题是真的妙啊，好像题解给得是O(n)的树形DP，点分治nlogn的1700ms过了。。。n再大一点就没了。

对于点分治，我们需要计算经过rt点所有路径的所有贡献，反过来计算，我们需要计算每一种颜色在经过rt点的路径有多少条

首先对于每种颜色cc的点，只需要考虑哪些离rt最近这种颜色点，来计算颜色cc在经过rt的路径中有多少条

假设这个点是t,它是rt的子树v中的一个点，且该点到rt的路径上没有c[t]这种颜色，那么ans+=(allson-son[v])*son[t]，也就是从son[t]出发到除了v子树的其他子树的所有点的路径，都会带上c[t]这种颜色。

used[cc]是表示在当前搜索遍历的情况下，cc的颜色出现了多少个点，calc(rt)的最开始就要把used[c[rt]]++

但是由于题目要求的路径是无序点对，我们rt的其他子树上已经计算过了c[t]这种颜色到v子树中的t子树中的这些点的路径，我们需要把他减去，减去用到了col[]数组和tmp[]数组配合，col[cc]表示cc这种颜色在已经遍历过的节点中的子树大小的累加和，tmp[cc]表示在当前遍历的子树中的cc这种颜色子树大小的累加和，这样就可以容斥掉了。

最后不要忘记used[c[rt]]--，计算c[rt]这个颜色的贡献那么就是直接用子树大小算就行了，注意也是无序点对的路径，

//有时间研究一下树形DP

 

#include<bits/stdc++.h>
using namespace std;

const int maxl=2e5+10;

int n,cnt,cas;
long long ans,allson,sum;
int ehead[maxl],c[maxl],used[maxl],s[maxl];
long long tmp[maxl],col[maxl];
vector<int> us1,us2;
struct ed
{
	int to,nxt;
}e[maxl<<1];
bool vis[maxl];
struct centertree
{
	int n,ans,mi;
	int son[maxl];
	inline void dfs(int u,int fa)
	{
		int v,mx=0;son[u]=1;
		for(int i=ehead[u];i;i=e[i].nxt)
		{
			v=e[i].to;
			if(vis[v] || v==fa) continue;
			dfs(v,u);
			son[u]+=son[v];
			mx=max(mx,son[v]);
		}
		mx=max(mx,n-son[u]);
		if(mx<mi)
		{
			ans=u;
			mi=mx;
		}
	}
	inline int getcenter(int u)
	{
		ans=0;mi=n+1;
		dfs(u,0);
		return ans;
	}
}tree;

inline void add(int u,int v)
{
	e[++cnt].to=v;e[cnt].nxt=ehead[u];ehead[u]=cnt;
}

inline void prework()
{
	for(int i=1;i<=n;i++)
	{
		scanf("%d",&c[i]);
		ehead[i]=0,vis[i]=false;
	}
	int u,v;cnt=0;
	for(int i=1;i<n;i++)
	{
		scanf("%d%d",&u,&v);
		add(u,v);add(v,u);
	}
}

inline int calcson(int u,int fa)
{
	int ret=1,v;
	for(int i=ehead[u];i;i=e[i].nxt)
	{
		v=e[i].to;
		if(vis[v] || v==fa) continue;
		ret+=calcson(v,u);
	}
	return ret;
}

inline int dfs(int u,int fa)
{
	used[c[u]]++;
	int v,ret=1;
	for(int i=ehead[u];i;i=e[i].nxt)
	{
		v=e[i].to;
		if(v==fa || vis[v]) continue;
		ret+=dfs(v,u);
	}
	used[c[u]]--;
	if(used[c[u]]==0)
	{
		ans+=ret*allson,tmp[c[u]]+=ret;
		us2.push_back(c[u]);
		us1.push_back(c[u]);
	}
	return ret;
}

inline void calc(int u)
{
	used[c[u]]++;
	sum=allson=1;
	us1.clear();int top=0;
	int v;
	for(int i=ehead[u];i;i=e[i].nxt)
	{
		v=e[i].to;
		if(vis[v]) continue;
		s[++top]=calcson(v,u);
		allson+=s[top];
	}
	int id=0,ret;
	for(int i=ehead[u];i;i=e[i].nxt)
	{
		v=e[i].to;
		if(vis[v]) continue;
		allson-=s[++id];us2.clear();
		ret=dfs(v,u);
		for(int cc : us2)
		{
			ans-=tmp[cc]*col[cc];
			col[cc]+=tmp[cc];tmp[cc]=0;
		}
		allson+=s[id];
		ans+=sum*s[id];sum+=s[id];
	}
	for(int cc: us1)
		col[cc]=0;
	used[c[u]]--;
}

inline void solv(int u)
{
	vis[u]=true;
	calc(u);int v,rt;
	for(int i=ehead[u];i;i=e[i].nxt)
	{
		v=e[i].to;
		if(vis[v]) continue;
		tree.n=tree.son[v];
		rt=tree.getcenter(v);
		solv(rt);
	}
}

inline void mainwork()
{
	ans=0;tree.n=n;
	int rt=tree.getcenter(1);
	solv(rt);
}

inline void print()
{
	++cas;
	printf("Case #%d: %lld\n",cas,ans);
}

int main()
{
	while(~scanf("%d",&n))
	{
		prework();
		mainwork();
		print();
	}
	return 0;
}