【bzoj3522】[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，是可以O(n2)暴力的。
三个点在一条链上这种情况是不可能的，所以一定是以某一个点为中心的三棵不同子树上的三个点。那么可以枚举这个中心，dfs出每棵子树每个深度的节点有多少个，然后乘法原理统计答案就好了。
对于答案统计的部分如果不理解可以看看这篇博客：
http://blog.csdn.net/slongle_amazing/article/details/50799448

#include<bits/stdc++.h>using namespace std;inline int read(){    int x = 0, f = 1; char c = getchar();    while(!isdigit(c)) { if(c == '-') f = -1; c = getchar(); }    while(isdigit(c)) { x = x * 10 + c - '0'; c = getchar(); }    return x * f;}const int N = 5000 + 10, M = 10000 + 10;int n, tot, mx;int to[M], nxt[M], hd[N], cnt[N], f[N], g[N];long long ans;inline void insert(int u, int v){    to[++tot] = v; nxt[tot] = hd[u]; hd[u] = tot;    to[++tot] = u; nxt[tot] = hd[v]; hd[v] = tot;}void init(){    n = read();    for(int i = 1; i < n; i++) insert(read(), read());}void dfs(int u, int fa, int dep){    cnt[dep]++;    mx = max(mx, dep);    for(int i = hd[u]; i; i = nxt[i])        if(to[i] != fa)            dfs(to[i], u, dep+1);}void work(){    for(int k = 1; k <= n; k++){        memset(f, 0, sizeof(f));        memset(g, 0, sizeof(g));        for(int i = hd[k]; i; i = nxt[i]){            dfs(to[i], k, 1);            for(int j = 1; j <= mx; j++){                ans += f[j] * cnt[j];                f[j] += g[j] * cnt[j];                g[j] += cnt[j];            }            for(int j = 1; j <= mx; j++) cnt[j] = 0;        }    }    printf("%lld\n", ans);}int main(){    init();    work();    return 0;}