HDU-4625 JZPTREE (树上dp，第二类斯特林数)

题意

给出$n$个结点的一棵树，对每一个点$x$求所有点到$x$的距离的$k$次方之和。$1leq nleq 50000, 1leq kleq 500$。

思路

用$Tree_x$表示这棵树以$x$为根，$f(x,k)$表示所有点到$x$的距离的$k$次方之和，$dis(x,y)$表示结点$x$和$y$之间的距离，题意即求：
$$
f(x,k) = sum_{yin Tree_x}{dis(x,y)^k} ag{1}
$$
我们要将之转换为能在树上做递推的递推式，并且时间复杂度在可接受的$O(nk)$范围之内。

《具体数学》中有这样的等式：
$$
n^k=sum_{i=1}^k{nchoose i} imes S_2(k,i) imes i! ag{2}
$$
其中$S_2(n,k)$为第二类斯特林数，表示把$n$个不同的球放到$k$个相同的盒子里，并且盒子不能为空，有多少种方案数，显然有：
$$
S_2(n,k)=S_2(n-1,k-1)+S_2(n-1,k) imes k ag{3}
$$
式子 (2) 的意义是：把$k$个不同的球放到$n$个不同的盒子里，显然有$n^k$种方法。或者先决定放到哪$1leq i leq k$个盒子里，再乘以$S_2(k,i)$就表示把$k$个不同的球放到这$i$个相同的盒子里的方案数，最后乘以阶乘表示放到不同的盒子里。

结合式子 (1) (2)，有：
$$
egin{split}
f(x,k) &= sum_{yin Tree_x}{dis(x,y)^k} \
&= sum_{yin Tree_x}sum_{i=1}^k{dis(x,y)choose i} imes S_2(k,i) imes i! \
&= sum_{i=1}^k S_2(k,i) imes i! imes sum_{yin Tree_x}{dis(x,y)choose i}
end{split} ag{4}
$$
记$son_x$表示与$x$邻接的结点，令：
$$
egin{split}
dp(x,i) &=& sum_{yin Tree_x}{dis(x,y)choose i} \
&=& sum_{yin son_x}sum_{zin Tree_y}{dis(y,z)+1choose i}
end{split} ag{5}
$$
而组合数有如下的帕斯卡等式：
$$
{nchoose k} = {n-1 choose k} + {n-1choose k-1} ag{6}
$$
结合式子 (5) (6) ，有
$$
egin{split}
dp(x,i) &=& sum_{yin son_x}sum_{zin Tree_y}left[ {dis(y,z)choose i} + {dis(y,z)choose i-1} ight] \
&=& sum_{yin son_x}[dp(y,i)+dp(y,i-1)]
end{split} ag{7}
$$
这就是一个可以在树上递推的递推式，而且能在$O(nk)$的时间复杂度内处理出所有的$dp(x,i)$。

结合式子 (4) (5)，有
$$
f(x,k) = sum_{i=1}^k S_2(k,i) imes i! imes dp(x,i) ag{8}
$$
因为$S_2(k,i) imes i!$是可以$O(k^2)$预处理的，所以总的时间复杂度为$O(k^2+nk)$。

(1)式是基于以$x$为树根的，代码实现中我们不能把每个结点都作为树根跑$n$次dfs，这样时间复杂度就不是$O(nk)$了。固定一个根，一次dfs处理每个结点子树部分，再跑一次dfs处理父亲结点那个分支即可。

代码实现

#include <cstdio>
#include <cstring>
const int N = 50010, K = 510, mod = 10007;
struct Edge
{
    int to, nex;
} edge[N<<1];
int fac[K], s[K][K], dp[N][K], head[N];
int cnt_e;
void add_edge(int u, int v) {
    edge[++cnt_e].to = v;
    edge[cnt_e].nex = head[u];
    head[u] = cnt_e;
}
void init1() {
    for (int i = fac[0] = 1; i < K; i++) fac[i] = fac[i-1] * i % mod;
    for (int i = s[0][0] = 1; i < K; i++) for (int j = 1; j <= i; j++) s[i][j] = (s[i-1][j-1] + j * s[i-1][j]) % mod;
}
void init2() {
    cnt_e = 0;
    memset(head, 0, sizeof(head));
}
void dfs1(int x, int fa, int k) {
    memset(dp[x], 0, sizeof(int) * (k + 1));
    dp[x][0] = 1;
    for (int i = head[x], y; y = edge[i].to, i; i = edge[i].nex) {
        if (y ^ fa) {
            dfs1(y, x, k);
            dp[x][0] += dp[y][0];
            for (int j = 1; j <= k; j++) dp[x][j] += dp[y][j] + dp[y][j-1];
        }
    }
    for (int i = 1; i <= k; i++) dp[x][i] >= mod ? dp[x][i] %= mod : 0;
}
int fa_dp(int x, int y, int i) {
    return dp[x][i] - dp[y][i] - (i > 0 ? dp[y][i-1] : 0) + 2 * mod;
}
void dfs2(int x, int fa, int k) {
    for (int i = head[x], y; y = edge[i].to, i; i = edge[i].nex) {
        if (y ^ fa) {
            for (int j = k; j > 0; j--) (dp[y][j] += fa_dp(x, y, j) + fa_dp(x, y, j - 1)) %= mod;
            dp[y][0] = dp[x][0];
            dfs2(y, x, k);
        }
    }
}
int get_ans(int x, int k) {
    long long res = 0ll;
    for (int i = 1; i <= k; i++) res += 1ll * fac[i] * s[k][i] * dp[x][i];
    return res % mod;
}

int main() {
    init1();
    int T;
    scanf("%d", &T);
    while (T--) {
        init2();
        int n, k;
        scanf("%d %d", &n, &k);
        for (int i = 0, u, v; i < n - 1; i++) {
            scanf("%d %d", &u, &v);
            add_edge(u, v), add_edge(v, u);
        }
        dfs1(1, 0, k);
        dfs2(1, 0, k);
        for (int i = 1; i <= n; i++) printf("%d
", get_ans(i, k));
    }
    return 0;
}

View Code