BZOJ3522[Poi2014]Hotel——树形DP

题目描述

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

输入

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

输出

让吉丽满意的方案数。

样例输入

7
1 2
5 7
2 5
2 3
5 6
4 5

样例输出

5

提示

【样例解释】

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

【数据范围】

n≤5000

数据范围比较小，考虑O(n²)树形DP。

满足要求的情况一定是一个点往外连出三条链，这三条链的端点就是要选的点。

如果把中间那个点当做根，那么这三个点就分别是根节点3个子树上.

那么我们不妨枚举根节点,对于每个根节点枚举子树统计答案.DP方程是f[i][j]=f[i][j-1]*s[i].

其中f[i][j]表示以当前枚举的点为根时已经选了j个深度为i的点(其中1<=j<=3)，s[i]则表示当前枚举的根的子树中深度为i的点的个数。

dfs完根的每棵子树转移时枚举i，j转移。但要注意j要倒序枚举，防止前面状态影响后面状态。

每dfs完一棵子树要清空s数组。最后对于每个点为根节点的答案加和就好了。

#include<queue>
#include<cmath>
#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
#define ll long long
using namespace std;
int n;
int x,y;
ll f[5010][4];
int head[5010];
int to[100010];
int next[100010];
int tot;
ll ans;
int s[50010];
int d[50010];
void add(int x,int y)
{
    tot++;
    next[tot]=head[x];
    head[x]=tot;
    to[tot]=y;
}
void dfs(int x,int fa)
{
    for(int i=head[x];i;i=next[i])
    {
        if(to[i]!=fa)
        {
            d[to[i]]=d[x]+1;
            s[d[to[i]]]++;
            dfs(to[i],x);
        }
    }
}
int main()
{
    scanf("%d",&n);
    for(int i=1;i<n;i++)
    {
        scanf("%d%d",&x,&y);
        add(x,y);
        add(y,x);
    }
    for(int i=1;i<=n;i++)
    {
        memset(d,0,sizeof(d));
        memset(f,0,sizeof(f));
        for(int j=1;j<=n;j++)
        {
            f[j][0]=1;
        }
        for(int j=head[i];j;j=next[j])
        {
            memset(s,0,sizeof(s));
            d[to[j]]=1;
            s[1]++;
            dfs(to[j],i);
            for(int k=3;k>=1;k--)
            {
                for(int l=1;s[l];l++)
                {
                    f[l][k]+=f[l][k-1]*s[l];
                }
            }
        }
        for(int j=1;f[j][3];j++)
        {
            ans+=f[j][3];
        }
    }
    printf("%lld",ans);
}