树形DP 统计树中长度为K的路径数量——Distance in Tree

一、问题描述

给出一棵n个节点的树，统计树中长度为k的路径的条数（1<=n<=50000 , 1<=k<=500）。

二、解题思路

设d[i][k]表示以i为根节点长度为k的路径数目

三、代码实现

#include<stdio.h>
#include<iostream>
#include<algorithm>
#include<cstring>
#include<vector>
using namespace std;

const int maxn = 50000 + 10;
const int maxk = 500 + 10;
vector<int>G[maxn];
int d[maxn][maxk],ans;
int n;

void init()
{
    ans = 0;
    memset(d, 0, sizeof(d));
    for (int i = 0; i <= n; i++)    G[i].clear();
}

void dfs(int root, int fa, int k)
{
    d[root][0] = 1;
    for (int i = 0; i < (int)G[root].size(); i++)
    {
        int son = G[root][i];
        if (son == fa)  continue;        //无向图，注意保持dfs的方向
        dfs(son, root, k);
        for (int j = 0; j < k; j++)  ans += d[son][j] * d[root][k - 1 - j];
        for (int j = 1; j <= k; j++)    d[root][j] += d[son][j - 1];
    }
}


int main()
{
    int k;
    while (scanf("%d%d",&n,&k) == 2)
    {
        int a, b;
        init();
        for (int i = 0; i < n - 1; i++)
        {
            scanf("%d%d", &a, &b);
            G[a].push_back(b);
            G[b].push_back(a);
        }
        dfs(1, 0, k);
        printf("%d
", ans);
    }
    return 0;
}