[bzoj3522] [Poi2014]Hotel

　　枚举一个点rt作为三个房间的中点。

　　那三个房间肯定在rt的不同子树内，且深度相同。

　　f1[i],f2[i],f3[i]分别表示深度为i，取1,2,3个点的方案数。

　　时间复杂度O(n^2)

#include<cstdio>
#include<iostream>
#include<cstring>
#include<algorithm>
#define ll long long
using namespace std;
const int maxn=5023;
struct zs{int too,pre,dis;}e[maxn<<1];int tot,last[maxn];
ll f3[maxn],ans;
int f2[maxn],f1[maxn],s[maxn],t,st[maxn],top;
int i,j,k,n,m;

int ra;char rx;
inline int read(){
rx=getchar(),ra=0;
while(rx<'0'||rx>'9')rx=getchar();
while(rx>='0'&&rx<='9')ra*=10,ra+=rx-48,rx=getchar();return ra;
}
inline void insert(int a,int b){
e[++tot].too=b,e[tot].pre=last[a],last[a]=tot,
e[++tot].too=a,e[tot].pre=last[b],last[b]=tot;
}
void dfs(int x,int fa,int dis){
s[++t]=dis;
for(int i=last[x];i;i=e[i].pre)if(e[i].too!=fa)dfs(e[i].too,x,dis+1);
}
int main(){
n=read();register int k;
for(i=1;i<n;i++)insert(read(),read());
for(i=1;i<=n;i++){
top=0;
for(j=last[i];j;j=e[j].pre)st[++top]=e[j].too;
if(top<3)continue;
for(j=1;j<=top;j++){
t=0,dfs(st[j],i,1);
for(k=1;k<=t;k++)f3[s[k]]+=f2[s[k]];
for(k=1;k<=t;k++)f2[s[k]]+=f1[s[k]];
for(k=1;k<=t;k++)f1[s[k]]++;
}
for(j=1;j<n;j++)ans+=f3[j];
memset(f1,0,n<<2),memset(f2,0,n<<2),memset(f3,0,n<<3);
}
printf("%lld\n",ans);
}

View Code