bzoj3522 [Poi2014]Hotel dfs（DP）

其实这题勉强可以算是个DP= =
一开始想的是设f[i][j]表示以i为根，距离为j的答案，然后发现转移好像很复杂，写出来了但是MLE= =。。懒得开滚动了。
然后发现其实并不用这么复杂= =
%题解：
首先可以发现，满足条件的点对一定是“有一个中心点，三个点到中心点的距离相等，且三个点分别在不同子树中”这种情况。
那么枚举中心点，然后遍历子树，统计答案。
cnt[i]为临时数组，表示当前子树中深度为i的点有多少个。
f[i]表示，当前点已经遍历过的子树中，深度为i的点有多少个。
g[i]表示，当前点已经遍历过的子树中，两个点深度都为i的点对有多少个。
那么ans = ∑(g[i] * cnt[i])。

#include<cstdio>
#include<algorithm>
#include<cstring>
#define fo(i,a,b) for(int i=a;i<=b;i++)
#define fd(i,a,b) for(int i=a;i>=b;i--)
using namespace std;
const int N=5e3+5;
typedef long long ll;
ll ans,f[N],g[N];
int cnt[N];
int n,in[N],tot;
int head[N<<1],go[N<<1],next[N<<1],fa[N];
inline void add(int x,int y)
{
    go[++tot]=y;
    next[tot]=head[x];
    head[x]=tot; 
}
inline void dfs(int x,int fa,int dep)
{
    cnt[dep]++;
    for(int i=head[x];i;i=next[i])
    {
        int v=go[i];
        if (v!=fa)
        {
            dfs(v,x,dep+1);
        }
    }
}
int main()
{
    scanf("%d",&n);
    fo(i,1,n-1)
    {
        int x,y;
        scanf("%d%d",&x,&y);
        add(x,y);
        add(y,x); 
        in[y]++;
    }
    fo(x,1,n)
    {
        memset(f,0,sizeof(f));
        memset(g,0,sizeof(g));
        for(int i=head[x];i;i=next[i])
        {
            int v=go[i];
            memset(cnt,0,sizeof(cnt));
            dfs(v,x,1);
            fo(j,1,n)
            {
                ans+=1ll*g[j]*cnt[j];
                g[j]+=f[j]*cnt[j];
                f[j]+=cnt[j];
            }
        }
    }
    printf("%lld\n",ans);
}