BZOJ 3522 Poi2014 Hotel DFS

题目大意：给定一棵树，求有多少无序三元组(x,y,z)满足x,y,z互不相等且Dis(x,y)=Dis(y,z)=Dis(x,z)

三个点在树上有两种情况

第一种是三点共链 第二种是存在且仅存在一个中心点满足三个点分别在这个点的三个不同的出边的方向

第一种情况显然无解

第二种情况一定满足三个点到中心点的距离相等

由于n<=5000因此直接枚举中心点然后枚举中心点的每一条出边DFS即可

时间复杂度O(n^2)

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
#define M 5050
using namespace std;
struct abcd{
	int to,next;
}table[M<<1];
int head[M],tot;
int n;
long long ans;
int temp[M],f[M],g[M];
void Add(int x,int y)
{
	table[++tot].to=y;
	table[tot].next=head[x];
	head[x]=tot;
}
void DFS(int x,int from,int dpt)
{
	int i;
	temp[dpt]++;
	for(i=head[x];i;i=table[i].next)
		if(table[i].to!=from)
			DFS(table[i].to,x,dpt+1);
}
int main()
{
	int i,j,x,y;
	cin>>n;
	for(i=1;i<n;i++)
	{
		scanf("%d%d",&x,&y);
		Add(x,y);Add(y,x);
	}
	for(x=1;x<=n;x++)
	{
		memset(f,0,sizeof f);
		memset(g,0,sizeof g);
		for(i=head[x];i;i=table[i].next)
		{
			memset(temp,0,sizeof temp);
			DFS(table[i].to,x,1);
			for(j=1;j<=n;j++)
			{
				ans+=(long long)g[j]*temp[j];
				g[j]+=f[j]*temp[j];
				f[j]+=temp[j];
			}
		}
	}
	cout<<ans<<endl;
	return 0;
}