E67 树形DP P3565 [POI2014] HOT-Hotels

// 树形DP O(n^2)
#include <iostream>
#include <cstring>
#include <algorithm>
using namespace std;

typedef long long LL;
const int N=5005;
int head[N],idx;
struct E{int to,ne;}e[N<<1];
void add(int x,int y){
  e[++idx]={y,head[x]},head[x]=idx;
}
int n,mxd;
LL d[N],f[N],g[N],ans;

void dfs(int u,int fa,int dep){
  mxd=max(mxd,dep); //最大深度
  d[dep]++;         //深dep的点数
  for(int i=head[u];i;i=e[i].ne){
    int v=e[i].to;
    if(v==fa)continue;
    dfs(v,u,dep+1);
  }
}
int main(){
  scanf("%d",&n);
  for(int i=1,x,y;i<n;i++){
    scanf("%d%d",&x,&y);
    add(x,y),add(y,x);
  }
  for(int i=1;i<=n;i++){
    memset(f,0,sizeof f);
    memset(g,0,sizeof g);
    for(int j=head[i];j;j=e[j].ne){
      memset(d,0,sizeof(d)); mxd=0;
      dfs(e[j].to,i,1);
      for(int k=1;k<=mxd;k++){
        ans+=f[k]*d[k];  //选三个点
        f[k]+=g[k]*d[k]; //选两个点
        g[k]+=d[k];      //选一个点
      }
    }
  }
  printf("%lld\n",ans);
}