bzoj 3522: Hotel dfs

题目大意

在无边权树上求三个点,使两两点的距离等。求方案数\((n\leq 5000)\)

题解

我们知道三个点在树上的关系只有两种

• 三点共链
• 三点不共连
  (这不是废话吗)
  我们发现三点共链肯定不满足条件,于是我们知道这三点一定不在同一条链上
  又由于这个条件的特性
  我们知道,对着这三个点来说肯定有一个中心点,就是到三个点的距离相等的点
  (这三个点关于中心点称菊花状放射分布)
  所以我们枚举这个中心点即可
  \(O(n^2)\)

Code

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
typedef long long ll;
inline void read(int &x){
	x=0;char ch;bool flag = false;
	while(ch=getchar(),ch<'!');if(ch == '-') ch=getchar(),flag = true;
	while(x=10*x+ch-'0',ch=getchar(),ch>'!');if(flag) x=-x;
}
inline int cat_max(const int &a,const int &b){return a>b ? a:b;}
inline int cat_min(const int &a,const int &b){return a<b ? a:b;}
const int maxn = 5010;
struct Node{
	int to,next;
}G[maxn<<1];
int head[maxn],cnt;
void add(int u,int v){
	G[++cnt].to = v;
	G[cnt].next = head[u];
	head[u] = cnt;
}
inline void insert(int u,int v){
	add(u,v);add(v,u);
}
int num[maxn],lim,t1[maxn],t2[maxn];
#define v G[i].to
void dfs(int u,int fa,int dis){
	++num[dis];
	lim = max(lim,dis);
	for(int i = head[u];i;i=G[i].next){
		if(v == fa) continue;
		dfs(v,u,dis+1);
	}
}
#undef v
int main(){
	int n;read(n);
	for(int i=1,u,v;i<n;++i){
		read(u);read(v);
		insert(u,v);
	}
	ll ans = 0;
	for(int u=1;u<=n;++u){
		memset(t1,0,sizeof t1);
		memset(t2,0,sizeof t2);
		lim = 0;
		for(int i = head[u];i;i=G[i].next){
			memset(num,0,sizeof num);
			dfs(G[i].to,u,1);
			for(int j=1;j<=lim;++j){
				ans += t2[j]*num[j];
				t2[j] += t1[j]*num[j];
				t1[j] += num[j];
			}
		}
	}
	printf("%lld\n",ans);
	getchar();getchar();
	return 0;
}

不过好像还有更好的姿势可以做到\(O(n)\)...