题解 CF161D Distance in Tree

设 \(f_{i,k}\) 表示 \(i\) 子树内和 \(i\) 距离为 \(k\) 的点数。

\[f_{i,0}=1 \\ f_{i,j}=\sum_{v\in son(i)}f_{v,j-1} \]

答案为：

\[\sum_{u,v\in son(p),u\neq v} \sum_{i=0}^k f_{u,i}\times f_{v,k-i} \\ \]

考虑一个点 \(v\) 的贡献：

\[\sum_{u\in son(p)}\sum_{i=0}^kf_{u,i}\times\sum_{v\in son(p),v\neq u}f_{v,k-i} \]

设：

\[sum_{u,i}=\sum_{j\in son(u)} f_{j,i} \]

则答案为：

\[\sum_{u\in son(p)}\sum_{i=0}^k f_{u,i}\times (sum_{p,k-i}-f_{u,k-i}) \]

时间复杂度 \(\mathcal{O}(nk)\)，点分治可以做到 \(\mathcal{O}(n\log_2 n)\)，不过本题没有必要。

代码如下：

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int N = 5e4 + 5 ,K = 505 ,M = 1e5 + 5;
struct Edge {
    int to ,next;
    Edge (int to = 0 ,int next = 0) : to(to) ,next(next) {}
}G[M]; int head[N] ,idx;
inline void add(int u ,int v) {
    G[++idx] = Edge(v ,head[u]); head[u] = idx;
    G[++idx] = Edge(u ,head[v]); head[v] = idx;
}
int f[N][K] ,sum[N][K];
LL ans;
int k;
inline void rcjx_dp(int now ,int fa) {
    f[now][0] = 1;
    for (int i = head[now]; i ; i = G[i].next) {
        int v = G[i].to;
        if (v == fa) continue;
        rcjx_dp(v ,now);
        for (int j = 0; j <= k; j++) {
            sum[now][j] += f[v][j];
            if (j < k) f[now][j + 1] += f[v][j];
        }
    }
    ans += f[now][k] * 2;
    for (int i = head[now]; i ; i = G[i].next) {
        int v = G[i].to;
        if (v == fa) continue;
        for (int j = 0; j <= k - 2; j++)
            ans += 1ll * f[v][j] * (sum[now][k - j - 2] - f[v][k - j - 2]);
    }
}
int n;
signed main() {
    scanf("%d%d" ,&n ,&k);
    for (int i = 1; i <= n - 1; i++) {
        int u ,v;
        scanf("%d%d" ,&u ,&v);
        add(u ,v);
    }
    rcjx_dp(1 ,0);
    printf("%lld\n" ,ans / 2);
    return 0;
}