cf161 D. Distance in Tree（树形dp）

题意：

求一棵树中距离为k的点对数量。点对没有顺序，即 (a,b) 和 (b,a) 是一样的。

n<=5e4, k<=500

思路：

dfs1处理出 cnt[u][d] ，表示以u为根节点的子树中与u的距离为d的点的数量，距离大于k的点不用考虑。

dfs2计算答案。经过节点u且长为k的路径有两种，第一种是以u为端点，数量是 cnt[u][k]；

第二种路径是以u为中间点，分成长为 d 和 d2 的两段（d和d2均不小于1且和为k）。在u的某棵子树v中取 到u的距离为d的点 作为路径起点，有 cnt[v][d-1] 个；然后在除v之外的u的其他子树中取终点，有 cnt[u][d2]-cnt[v][d2-1] 个，两者相乘并加到答案中。如此计算u的所有子树。

注意第二种路径会计算两次，需除以2。

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const int N = 5e4 + 5, M = N * 2;
int h[N], e[M], ne[M], idx;
void add(int a, int b)
{
    e[++idx] = b, ne[idx] = h[a], h[a] = idx;
}

int n, k, cnt[N][505];
void dfs1(int u, int fa)
{
    cnt[u][0] = 1;
    for(int i = h[u]; i; i = ne[i])
    {
        int v = e[i]; if(v == fa) break;
        dfs1(v, u);

        for(int d = 1; d <= k; d++)
            cnt[u][d] += cnt[v][d-1];
    }
}

ll ans;
void dfs2(int u, int fa)
{
    ans += cnt[u][k];
    ll res = 0;
    for(int i = h[u]; i; i = ne[i])
    {
        int v = e[i]; if(v == fa) break;
        for(int d = 1; d <= k-1; d++)
        {
            int d2 = k - d;
            res += (ll)(cnt[v][d-1]) * (cnt[u][d2] - cnt[v][d2-1]);
        }

        dfs2(v, u);
    }
    ans += res / 2;
}

signed main()
{
    scanf("%d%d", &n, &k);
    for(int i = 1; i < n; i++)
    {
        int a, b; scanf("%d%d", &a, &b);
        add(a, b), add(b, a);
    }

    dfs1(1, -1), dfs2(1, -1);

    cout << ans;

    return 0;
}