HDU 4625 JZPTREE（树形DP+组合数学）

Description

给出一棵$n$个节点的树，边权为一，求$\sum\limits_{u,v}dist(u,v)^k$

Input

第一行一整数$T$表示用例组数，每组用例首先输入两个整数$n,k$，之后$n-1$行每行两个整数表示一条树边

$(1\le T\le 5,1\le n\le 50000,1\le k\le 500)$

Output

输出$\sum\limits_{u,v}dist(u,v)^k$，结果模$10007$

Sample Input

1
5 3
1 2
2 3
3 4
4 5

Sample Output

100
37
18
37
100

Solution

简记$d_{u,v}=dis(u,v)$，首先给出恒等式$n^k=\sum\limits_{i=1}^kA_n^i\cdot S(k,i)$，其中$n^k$可以表示将$k$个不同的小球放入$n$个不同的盒子的方案数，而这个方案数就等于从$n$个盒子选$i$个盒子非空，然后将$k$个不同小球放入这$i$个不同的盒子里的方案数，后者即为第二类斯特林数$S(k,i)$，故可以对答案进行化简

$$\sum\limits_{v=1}^nd_{u,v}^k=\sum\limits_{v=1}^n\sum\limits_{i=1}^kS(k,i)\cdot i!\cdot C_{d_{u,v}}^i=\sum\limits_{i=1}^kS(k,i)\cdot i!\sum\limits_{v=1}^nC_{d_{u,v}}^i$$
记 $dp[u][i]=\sum\limits_{v=1}^nC_{d_{u,v}}^i$，只要求出 $dp[u][i]$即可 $O(k)$得到固定 $u$点的答案

不妨以$1$为树根，记 $t(u)$表示以 $u$为根节点的子树，首先可以一遍树形$DP$得到 $f[u][i]=\sum\limits_{v\in t(u)}C_{d_{u,v}}^i$，这是由于经典等式 $C_n^m=C_{n-1}^m+C_{n-1}^{m-1}$，且对于 $u$的儿子节点$v$的子树中一点 $w$有$d_{u,w}=d_{v,w}+1$，故 $C_{d_{u,w}}^i=C_{d_{v,w}}^i+C_{d_{v,w}}^{i-1}$，进而有

$$f[u][i]=C_{d_{u,u}}^i+\sum\limits_{v\in son(u)}(f[v][i]+f[v][i-1])$$
注意到 $dp[1][i]=f[1][i]$，之后再一遍树形 $DP$，在已知 $u$节点父亲节点的$dp[fa][i]$后来求 $dp[u][i]$，注意到对于 $u$子树中一点$v$有 $d_{fa,v}=d_{u,v}+1$，而对于 $u$子树外一点$v$有 $d_{fa,v}=d_{u,v}-1$，故有
$$\begin{array}{rcl} dp[u][i]&=&\sum\limits_{v=1}^nC_{d_{u,v}}^i\\ &=&\sum\limits_{v\in t(u)}C_{d_{fa,v}-1}^i+\sum\limits_{v\not\in t(u)}C_{d_{fa,v}+1}^i\\ &=&\sum\limits_{v\in t(u)}C_{d_{fa,v}}^i-\sum\limits_{v\in t(u)}C_{d_{fa,v}-1}^{i-1}+\sum\limits_{v\not\in t(u)}C_{d_{fa,v}}^i+\sum\limits_{v\not\in t(u)}C_{d_{fa,v}}^{i-1}\\\ &=&dp[fa][i]+dp[fa][i-1]-\sum\limits_{v\in t(u)}C_{d_{fa,v}-1}^{i-1}-\sum\limits_{v\in t(u)}C_{d_{fa,v}}^{i-1}\\ &=&dp[fa][i]+dp[fa][i-1]-\sum\limits_{v\in t(u)}C_{d_{fa,v}-1}^{i-1}-\sum\limits_{v\in t(u)}C_{d_{fa,v}-1}^{i-1}-\sum\limits_{v\in t(u)}C_{d_{fa,v}-1}^{i-2}\\ &=&dp[fa][i]+dp[fa][i-1]-2f[u][i-1]-f[u][i-2] \end{array}$$
总时间复杂度 $O(Tnk)$

Code

#include<cstdio>
#include<vector>
using namespace std;
#define maxn 50005
#define mod 10007
int mul(int x,int y)
{
    int z=x*y;
    return z-z/mod*mod;
}
int add(int x,int y)
{
    x+=y;
    if(x>=mod)x-=mod;
    return x;
}
int T,n,k,dp[maxn][505],S[505][505],fact[505];
vector<int>g[maxn];
void init(int n=500)
{
    S[0][0]=1;
    for(int i=1;i<=n;i++)
    {
        S[i][i]=1;
        for(int j=1;j<i;j++)
            S[i][j]=add(S[i-1][j-1],mul(S[i-1][j],j));
    }
    fact[0]=1;
    for(int i=1;i<=n;i++)fact[i]=mul(i,fact[i-1]);
}
void dfs1(int u,int fa)
{
    dp[u][0]=1;
    for(int i=1;i<=k;i++)dp[u][i]=0;
    for(int i=0;i<g[u].size();i++)
    {
        int v=g[u][i];
        if(v==fa)continue;
        dfs1(v,u);
        dp[u][0]=add(dp[u][0],dp[v][0]);
        for(int j=1;j<=k;j++)
            dp[u][j]=add(dp[u][j],add(dp[v][j],dp[v][j-1]));
    }
}
void dfs2(int u,int fa)
{
    if(fa)
    {
        for(int j=k;j>=0;j--)
        {
            dp[u][j]=dp[fa][j];
            if(j>=1)dp[u][j]=add(dp[u][j],add(dp[fa][j-1],mul(dp[u][j-1],mod-2)));
            if(j>=2)dp[u][j]=add(dp[u][j],mod-dp[u][j-2]);
        }
    }
    for(int i=0;i<g[u].size();i++)
    {
        int v=g[u][i];
        if(v==fa)continue;
        dfs2(v,u);
    }
}
int main()
{
    init();
    scanf("%d",&T);
    while(T--)
    {
        scanf("%d%d",&n,&k);
        for(int i=1;i<=n;i++)g[i].clear();
        for(int i=1;i<n;i++)
        {
            int u,v;
            scanf("%d%d",&u,&v);
            g[u].push_back(v),g[v].push_back(u);
        }
        dfs1(1,0);
        dfs2(1,0);
        for(int i=1;i<=n;i++)
        {
            int ans=0;
            for(int j=1;j<=k;j++)
                ans=add(ans,mul(mul(S[k][j],fact[j]),dp[i][j]));
            printf("%d\n",ans); 
        }
    }
    return 0;
}