【BZOJ3522】【BZOJ4543】【POI2014】Hotel 树形DP 长链剖分 启发式合并

题目大意

​　　给你一棵树，求有多少个组点满足\(x\neq y,x\neq z,y\neq z,dist_{x,y}=dist_{x,z}=dist_{y,z}\)

​　　\(1\leq n\leq 100000\)

题解

​　　问题转换为有多少个组点满足\(dist_{i,x}=dist_{i,y}=dist_{i,z}\)

​　　我们考虑树形DP

​　　\(f_{i,j}=\)以\(i\)为根的子树中与\(i\)的距离为\(j\)的节点数

​　　\(g_{i,j}=\)以\(i\)为根的子树外选择一个点\(s\)满足\(s\)到\(i\)的距离为\(j\)，能新增的的方案数

​　　若\(v\)是\(u\)的重儿子，则：\(f_{u,j}+=f_{v,j-1},g_{u,j}+=g_{v,j+1}\)，这样就可以由\(u\)的重儿子转移到\(u\)

​　　否则：\(g_{u,j}+=g_{v,{j+1}}+f_{v,j-1}\times f_{u,j},f_{u,j}+=f_{v,j-1}\)

​　　答案为\(\sum f_{x,j}\times g_{y,j}\)，其中\(x\)是\(y\)的兄弟

​　　可以用长链剖分辅助转移

​　　时间复杂度：\(O(n)\)

​　　gjs大爷的长链剖分讲解

代码

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cstdlib>
#include<ctime>
#include<utility>
using namespace std;
typedef long long ll;
typedef pair<int,int> pii;
struct list
{
	int v[200010];
	int t[200010];
	int h[100010];
	int n;
	void clear()
	{
		n=0;
		memset(h,0,sizeof h);
	}
	void add(int x,int y)
	{
		n++;
		v[n]=y;
		t[n]=h[x];
		h[x]=n;
	}
};
list l;
ll ans;
ll f[100010];
ll g[200010];
int d[100010];
int bg[100010];
int ed[100010];
int ch[100010];
int t[100010];
int w[100010];
int ti;
void dfs(int x,int fa)
{
	d[x]=1;
	ch[x]=0;
	int i;
	for(i=l.h[x];i;i=l.t[i])
		if(l.v[i]!=fa)
		{
			dfs(l.v[i],x);
			if(d[l.v[i]]+1>d[x])
			{
				d[x]=d[l.v[i]]+1;
				ch[x]=l.v[i];
			}
		}
}
void dfs2(int x,int fa,int top)
{
	t[x]=top;
	w[x]=++ti;
	if(x==top)
		bg[top]=ti;
	ed[top]=ti;
	if(ch[x])
		dfs2(ch[x],x,top);
	int i;
	for(i=l.h[x];i;i=l.t[i])
		if(l.v[i]!=ch[x]&&l.v[i]!=fa)
			dfs2(l.v[i],x,l.v[i]);
}
ll& getf(int x,int y)
{
	return f[w[x]+y];
}
ll& getg(int x,int y)
{
	return g[2*(w[t[x]]-1)+2*d[t[x]]-d[x]+1-y];
}
void solve(int x,int fa)
{
	if(ch[x])
		solve(ch[x],x);
	int i,j;
	for(i=l.h[x];i;i=l.t[i])
		if(l.v[i]!=fa&&l.v[i]!=ch[x])
		{
			int v=l.v[i];
			solve(v,x);
			for(j=0;j<d[v];j++)
				ans+=getf(v,j)*getg(x,j+1);
			for(j=1;j<d[v];j++)
				ans+=getg(v,j)*getf(x,j-1);
			for(j=0;j<d[v];j++)
				getg(x,j+1)+=getf(v,j)*getf(x,j+1);
			for(j=1;j<d[v];j++)
				getg(x,j-1)+=getg(v,j);
			for(j=0;j<d[v];j++)
				getf(x,j+1)+=getf(v,j);
		}
	ans+=getg(x,0);
	getf(x,0)++;
}
int main()
{
	int n;
	scanf("%d",&n);
	l.clear();
	memset(bg,0,sizeof bg);
	memset(ed,0,sizeof ed);
	memset(f,0,sizeof f);
	memset(g,0,sizeof g);
	memset(d,0,sizeof d);
	memset(ch,0,sizeof ch);
	memset(t,0,sizeof t);
	memset(w,0,sizeof w);
	ans=0;
	ti=0;
	int i,x,y;
	for(i=1;i<n;i++)
	{
		scanf("%d%d",&x,&y);
		l.add(x,y);
		l.add(y,x);
	}
	dfs(1,0);
	dfs2(1,0,1);
	solve(1,0);
	printf("%lld\n",ans);
	return 0;
}