【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
, cnt
, f
, g
;
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;
}