【BZOJ3522】[Poi2014]Hotel 树形DP

【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

题解：由于可以O(n^2)搞，很容易想到枚举中间点（树根），然后用树形DP搞一搞

发现三个妹子一定在根的三个不同儿子的子树中，所以我们可以遍历根的每个儿子，设s[i]表示其中一个儿子的子树中，深度为i的点有多少个，然后设f[i][j]表示在j个不同儿子的子树中，各选出一个深度为i的点的方案数，可以得到下面的式子

f[i][j]+=f[i][j-1]*s[i]

注意f[0][..]要赋成0，并且j要从3到1循环（01背包的思想）

#include <cstdio>
#include <cstring>
#include <iostream>
using namespace std;
const int maxn=5010;
int n,cnt;
long long ans;
int to[maxn<<1],next[maxn<<1],head[maxn],dep[maxn];
long long s[maxn],f[4][maxn];
void add(int a,int b)
{
    to[cnt]=b,next[cnt]=head[a],head[a]=cnt++;
}
void dfs(int x,int fa)
{
    int i;
    s[dep[x]]++;
    for(i=head[x];i!=-1;i=next[i])
    {
        if(to[i]==fa)   continue;
        dep[to[i]]=dep[x]+1,dfs(to[i],x);
    }
}
int main()
{
    scanf("%d",&n);
    int i,j,k,l,a,b;
    memset(head,-1,sizeof(head));
    for(i=1;i<n;i++)
    {
        scanf("%d%d",&a,&b);
        add(a,b),add(b,a);
    }
    for(i=1;i<=n;i++)
    {
        memset(f,0,sizeof(f));
        for(k=1;k<=n;k++)    f[0][k]=1;
        for(j=head[i];j!=-1;j=next[j])
        {
            memset(s,0,sizeof(s));
            dep[to[j]]=1,dfs(to[j],i);
            for(l=3;l>=1;l--)
                for(k=1;s[k];k++)
                    f[l][k]+=f[l-1][k]*s[k];
        }
        for(k=1;f[3][k];k++)    ans+=f[3][k];
    }
    printf("%lld",ans);
    return 0;
}