codeforces 161D - Distance in Tree（树形dp）

题目大意：

求出树上距离为k的点对有多少个。

思路分析：

dp[i][j] 表示 i 的子树中和 i 的距离为 j 的点数有多少个。注意dp[i] [0] 永远是1的。

然后在处理完一颗子树后，就把自身的dp 更新。

更新之前更新答案。

如果这颗子树到 i 有 x 个距离为j的。那么答案就要加上 dp[i] [ k-j-1] * x;

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#define maxn 50005
using namespace std;
typedef long long ll;
ll dp[maxn][505];

int head[maxn];
int to[maxn<<1];
int next[maxn<<1];
int tot;
int n,k;
void addedge(int u,int v)
{
    tot++;
    next[tot]=head[u];
    to[tot]=v;
    head[u]=tot;
}

bool vis[maxn];
ll ans=0;
void dfs(int x)
{
    dp[x][0]=1;
    for(int p=head[x];p;p=next[p])
    {
        if(vis[to[p]])continue;

        vis[to[p]]=true;
        dfs(to[p]);

        for(int i=0;i<k;i++)
        {
            ans+=dp[to[p]][k-i-1]*dp[x][i];
        }
        for(int i=0;i<=k;i++)
        {
            if(i)dp[x][i]+=dp[to[p]][i-1];
            else dp[x][i]=1;
        }
    }
}
int main()
{
    scanf("%d%d",&n,&k);
    for(int i=1;i<n;i++)
    {
        int l,r;
        scanf("%d%d",&l,&r);
        addedge(l,r);
        addedge(r,l);
    }

    memset(vis,false,sizeof vis);
    memset(dp,0,sizeof dp);
    vis[1]=true;
    dfs(1);
    printf("%I64d\n",ans);
    return 0;
}

codeforces 161D - Distance in Tree（树形dp）