CF-161D Distance in Tree

题意： 给出一棵树，边长度皆为1，问距离为k的点对有多少组，相同点对算一组。

解法： 设dp[u][i]为以u为根的子树中与u的距离为i的点有多少，DFS一遍即可。

#include <stdio.h>
#include <string.h>
#include <algorithm>
#include <math.h>
using namespace std;
#define REP(i,a,b) for(int i=(a); i<(b); i++)
#define clr(a,b) memset(a,b,sizeof(a))
typedef long long lld;
const int INF = ~0u>>1;

int dp[50001][505];
int n,k;
struct Edge {
    int v,next;
}g[100001];
int tot,head[50001];
lld cnt;

void add_edge(int u, int v) {
    g[tot].v = v;
    g[tot].next = head[u];
    head[u] = tot ++;
}

void dfs(int u, int fa) {
    dp[u][0] = 1;
    for(int p=head[u]; ~p; p=g[p].next) {
        int v = g[p].v;
        if(v == fa) continue;
        dfs(v,u);
        REP(i,0,k) dp[u][i+1] += dp[v][i];
    }
    for(int p=head[u]; ~p; p=g[p].next) {
        int v = g[p].v;
        if(v == fa) continue;
        REP(i,0,k-1) dp[u][i+1] -= dp[v][i];
        REP(i,0,k-1) cnt += (lld) dp[v][i] * dp[u][k-i-1];
        REP(i,0,k-1) dp[u][i+1] += dp[v][i];
    }
    cnt += (lld) dp[u][k] << 1;
}

int main() {
    scanf("%d%d", &n, &k);
    clr(head, -1);
    int a,b;
    REP(i,0,n-1) {
        scanf("%d%d", &a, &b);
        add_edge(a,b);
        add_edge(b,a);
    }
    dfs(1,-1);
    printf("%lld\n", cnt>>1);
    return 0;
}