题目链接: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_x}{dis(x,y)^k} \tag{1}
$$
我们要将之转换为能在树上做递推的递推式,并且时间复杂度在可接受的$O(nk)$范围之内。
《具体数学》中有这样的等式:
$$
n^k=\sum_{i=1}^k{n\choose i}\times S_2(k,i)\times i! \tag{2}
$$
其中$S_2(n,k)$为第二类斯特林数,表示把$n$个不同的球放到$k$个相同的盒子里,并且盒子不能为空,有多少种方案数,显然有:
$$
S_2(n,k)=S_2(n-1,k-1)+S_2(n-1,k)\times k \tag{3}
$$
式子 (2) 的意义是:把$k$个不同的球放到$n$个不同的盒子里,显然有$n^k$种方法。或者先决定放到哪$1\leq i \leq k$个盒子里,再乘以$S_2(k,i)$就表示把$k$个不同的球放到这$i$个相同的盒子里的方案数,最后乘以阶乘表示放到不同的盒子里。
结合式子 (1) (2),有:
$$
\begin{split}
f(x,k) &= \sum_{y\in Tree_x}{dis(x,y)^k} \\
&= \sum_{y\in Tree_x}\sum_{i=1}^k{dis(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_x}{dis(x,y)\choose i}
\end{split} \tag{4}
$$
记$son_x$表示与$x$邻接的结点,令:
$$
\begin{split}
dp(x,i) &=& \sum_{y\in Tree_x}{dis(x,y)\choose i} \\
&=& \sum_{y\in son_x}\sum_{z\in Tree_y}{dis(y,z)+1\choose i}
\end{split} \tag{5}
$$
而组合数有如下的帕斯卡等式:
$$
{n\choose k} = {n-1 \choose k} + {n-1\choose k-1} \tag{6}
$$
结合式子 (5) (6) ,有
$$
\begin{split}
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)]
\end{split} \tag{7}
$$
这就是一个可以在树上递推的递推式,而且能在$O(nk)$的时间复杂度内处理出所有的$dp(x,i)$。
结合式子 (4) (5),有
$$
f(x,k) = \sum_{i=1}^k S_2(k,i)\times i!\times dp(x,i) \tag{8}
$$
因为$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; }