【DFS】BZOJ3522-[Poi2014]Hotel

【题目大意】

给出一棵树，求三个节点使得它们两两之间的距离相等，问共有多少种可能性？

【思路】

显然，这三个节点是关于一个中心点对称地辐射出去的。

枚举中心点，往它的各个子树跑Dfs。tmp[i]表示当前子树深度为i的节点个数，p1[i]表示之前的子树中（不包括当前的子树），深度为i的节点的个数，p2[i]表示到当前子树之前，选取两个不同子树的节点的方案数。

对于当前这一棵子树的深度i，显然有ans+=tmp[i]*p2[i]。然后更新下次要用的p1[]、p2[]，p2[i]+=p1[i]*tmp[i]，p1[i]+=tmp[i]。

不要忘了long long。

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int MAXN=5000+50;
int n;
int p1[MAXN],p2[MAXN],tmp[MAXN];
vector<int> E[MAXN];

void dfs(int fr,int u,int dep)
{
    tmp[dep]++;
    for (int i=0;i<E[u].size();i++)
    {
        int v=E[u][i];
        if (v==fr) continue;
        dfs(u,v,dep+1);
    }
}

void solve()
{
    ll ans=0;
    for (int i=1;i<=n;i++)
    {
        memset(p1,0,sizeof(p2));
        memset(p2,0,sizeof(p1));
        for (int j=0;j<E[i].size();j++)
        {
            memset(tmp,0,sizeof(tmp));
            int to=E[i][j];
            dfs(i,to,1);
            for (int k=1;k<=n;k++)
            {
                ans+=(ll)p2[k]*tmp[k];
                p2[k]+=p1[k]*tmp[k];
                p1[k]+=tmp[k];
            }
        }
    }
    printf("%lld\n",ans);
}

void init()
{
    scanf("%d",&n);
    for (int i=0;i<n-1;i++)
    {
        int x,y;
        scanf("%d%d",&x,&y);
        E[x].push_back(y);
        E[y].push_back(x);
    }
}

int main()
{
    init();
    solve();
    return 0;    
}