CodeForces 161 D.Distance in Tree（树形DP）

Description

给出一棵$n$个节点的树，边权均为$1$，问树上距离为$k$的点对数量，$(u,v)$和$(v,u)$视作一种

Input

第一行两整数$n,k$，之后输入$n-1$条树边$(1\le n\le 50000,1\le k\le 500)$

Output

输出树上距离为$k$的点对数

Sample Input

5 2
1 2
2 3
3 4
2 5

Sample Output

4

Solution

$dp[u][j]$表示$u$的子树中距$u$距离为$j$的点的数量，那么在不考虑$u$的儿子时有$dp[u][0]=0$，之后逐个考虑$u$的儿子，假设当前考虑$u$的儿子$v$，那么两点之间路径经过$u$且距离$k$的点对数量就会增加$dp[u][j]\cdot dp[v][k-j+1]$，之后再将$dp[v][j-1]$累加到$dp[u][j]$中

Code

#include<cstdio>
#include<iostream>
#include<cstring>
#include<algorithm>
#include<cmath>
#include<vector>
#include<queue>
#include<map>
#include<set>
#include<ctime>
using namespace std;
typedef long long ll;
typedef pair<int,int>P;
const int INF=0x3f3f3f3f,maxn=50001;
int n,k,dp[maxn][501],ans;
vector<int>g[maxn];
void dfs(int u,int fa)
{
    dp[u][0]=1;
    for(int i=0;i<g[u].size();i++)
    {
        int v=g[u][i];
        if(v==fa)continue;
        dfs(v,u);
        for(int j=0;j<k;j++)ans+=dp[u][j]*dp[v][k-j-1];
        for(int j=1;j<=k;j++)dp[u][j]+=dp[v][j-1];
    }
}
int main()
{
    scanf("%d%d",&n,&k);
    for(int i=1;i<n;i++)
    {
        int u,v;
        scanf("%d%d",&u,&v);
        g[u].push_back(v),g[v].push_back(u);
    }
    dfs(1,0);
    printf("%d\n",ans);
    return 0;
}