BZOJ 3522: [Poi2014]Hotel

Description

有一个树形结构的宾馆，n个房间，n-1条无向边，每条边的长度相同，任意两个房间可以相互到达。吉丽要给他的三个妹子各开（一个）房（间）。三个妹子住的房间要互不相同（否则要打起来了），为了让吉丽满意，你需要让三个房间两两距离相同。
有多少种方案能让吉丽满意？

Input

第一行一个数n。
接下来n-1行，每行两个数x,y，表示x和y之间有一条边相连。

Output

让吉丽满意的方案数。

Sample Input

7

1 2

5 7

2 5

2 3

5 6

4 5

Sample Output

5

HINT

【样例解释】

{1,3,5},{2,4,6},{2,4,7},{2,6,7},{4,6,7}

【数据范围】

n≤5000

分析

三个点在树上有两种情况
第一种是三点共链 第二种是存在且仅存在一个中心点满足三个点分别在这个点的三个不同的出边的方向
第一种情况显然无解
第二种情况一定满足三个点到中心点的距离相等
由于n<=5000因此直接枚举中心点然后枚举中心点的每一条出边DFS即可

代码

#include <bits/stdc++.h>

#define N 5050

struct NOTE
{
    int to,next;
}e[N << 1];

int cnt;
int next[N];

int tot[N],g[N],f[N];

long long ans;

int read()
{
    int x = 0, f = 1;
    char ch = getchar();
    while (ch < '0' || ch > '9') {if (ch == '-') f = -1; ch = getchar();}
    while (ch >= '0' && ch <= '9') {x = x * 10 + ch - '0'; ch = getchar();}
    return x * f;
}

void add(int x,int y)
{
    e[++cnt] = (NOTE){y,next[x]}; next[x] = cnt;
    e[++cnt] = (NOTE){x,next[y]}; next[y] = cnt;
}

void dfs(int x,int f,int d)
{
    tot[d]++;
    for (int i = next[x]; i; i = e[i].next)
    {
        if (e[i].to == f)
            continue;
        dfs(e[i].to, x, d + 1);
    }
}

int main()
{
    int n = read();
    for (int i = 1; i < n; i++)
    {
        int x = read(), y = read();
        add(x,y);
    }
    for (int x = 1; x <= n; x++)
    {
        memset(f,0,sizeof(f));
        memset(g,0,sizeof(g));
        for (int i = next[x]; i; i = e[i].next)
        {
            memset(tot,0,sizeof(tot));
            dfs(e[i].to,x,1);
            for (int j = 1; j <= n; j++)
            {
                ans += (long long)g[j] * tot[j];
                g[j] += f[j] * tot[j];
                f[j] += tot[j];
            }
        }
    }
    printf("%lld\n",ans);
}