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

Posted 2022-09-01 kangkang-

题目链接：HDU-4625 JZPTREE

题意

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

---

思路

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

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

结合式子 (1) (2)，有：
$$
\beginsplit
f(x,k) &= \sum_y\in Tree_xdis(x,y)^k \\
&= \sum_y\in Tree_x\sum_i=1^kdis(x,y)\choose i\times S_2(k,i) \times i! \\
&= \sum_i=1^k S_2(k,i)\times i!\times \sum_y\in Tree_xdis(x,y)\choose i
\endsplit \tag4
$$
记$son_x$表示与$x$邻接的结点，令：
$$
\beginsplit
dp(x,i) &=& \sum_y\in Tree_xdis(x,y)\choose i \\
&=& \sum_y\in son_x\sum_z\in Tree_ydis(y,z)+1\choose i
\endsplit \tag5
$$
而组合数有如下的帕斯卡等式：
$$
n\choose k = n-1 \choose k + n-1\choose k-1 \tag6
$$
结合式子 (5) (6) ，有
$$
\beginsplit
dp(x,i) &=& \sum_y\in son_x\sum_z\in Tree_y\left[ dis(y,z)\choose i + dis(y,z)\choose i-1\right] \\
&=& \sum_y\in son_x[dp(y,i)+dp(y,i-1)]
\endsplit \tag7
$$
这就是一个可以在树上递推的递推式，而且能在$O(nk)$的时间复杂度内处理出所有的$dp(x,i)$。

结合式子 (4) (5)，有
$$
f(x,k) = \sum_i=1^k S_2(k,i)\times i!\times dp(x,i) \tag8
$$
因为$S_2(k,i)\times 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\n", get_ans(i, k));
    
    return 0;

View Code