CodeForces - 161D

Posted 2021-06-04 吃花椒的妙酱

题目大意：求树上距离为k的点对数

思路：树形dp

dp[i][j]表示到点i距离为j的点的数量，从叶子维护到根

有转移方程dp[x][i] = ∑dp[son][i-1]

对于答案贡献有 dp[x][i] * dp[son][k-i-1]，因为儿子到当前结点x有距离1，从而我们只要考虑每对儿子和父亲就可以计算出答案

每次更新完儿子的dp后，就马上算答案，避免贡献算重复了

#define _CRT_SECURE_NO_WARNINGS
#include <iostream>
#include <cstring>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <algorithm>
#include <map>
#include <vector>
#include <queue>
#include <set>
using namespace std;
#define _for(i,a,b) for(int i=(a) ;i<=(b) ;i++)
#define _rep(i,a,b) for(int i=(a) ;i>=(b) ;i--)
#define mst(v,s) memset(v,s,sizeof(v))
#define pb push_back
#define IOS ios::sync_with_stdio(false)
#define int long long
#define inf 0x3f3f3f3f
#define INF 0x3f3f3f3f3f3f3f3f
typedef long long ll;
ll ans = 0;
int dp[51000][510];
vector <int > G[51000];
int n, k;
void dfs(int x, int fa)
{
    dp[x][0] = 1;//当自己的点数为1，即自己
    for (int i = 0; i < G[x].size(); i++)
    {
        int y = G[x][i];
        if (y == fa) continue;
        dfs(y, x);
        //先更新答案
        for (int j = 0; j < k; j++)
            ans += dp[x][j] * dp[y][k - j - 1];
        //再利用儿子更新当前结点
        for (int j = 0; j < k; j++)
            dp[x][j + 1] += dp[y][j];
    }
}
signed main(void )
{
    //!!!
    //freopen("data.txt", "r", stdin);
    //!!!
    IOS;
   cin >> n >> k;
    _for(i, 1, n - 1)
    {
        int x, y; cin >> x >> y;
        G[x].pb(y); G[y].pb(x);
    }
    dfs(1, 0);
    cout << ans << endl;
}