Distance in Tree 树形dp练习（树中两点距离为k的数量板子）

Distance in Tree

题面翻译

题目大意

输入点数为$N$一棵树

求树上长度恰好为$K$的路径个数

输入格式

第一行两个数字$N,K$,如题意

接下来的$N-1$行中,每行两个整数$u,v$表示一条树边$(u,v)$

输出格式

一个整数$ans$,如题意

在$Codeforces$上提交时记得用$\%I64d$哦.QwQ

数据范围

$1\le n\le 50000$

$1\le k\le 500$

样例输入 #1

5 2
1 2
2 3
3 4
2 5

样例输出 #1

4

样例 #2

样例输入 #2

5 3
1 2
2 3
3 4
4 5

样例输出 #2

2

提示

In the first sample the pairs of vertexes at distance 2 from each other are (1, 3), (1, 5), (3, 5) and (2, 4).

思路：根据数据范围，我没可以设置一个二维数组dp[i][j], 表示以i为根节点，到i的距离为j时的方案数量，那么我们的方案数量s可以表示为s += dp[u][j] * dp[v][k - 1 - j], 我们只需要求到i的距离为k - 1即可，因为i本身也是一个节点，做一下树形dp即可

#include<bits/stdc++.h>

using namespace std;

typedef long long ll;
typedef pair<ll, ll>PII;
const int N = 2e5 + 10;
const int MOD = 998244353;
const int INF = 0X3F3F3F3F;
const int dx[] = {-1, 1, 0, 0, -1, -1, +1, +1};
const int dy[] = {0, 0, -1, 1, -1, +1, -1, +1};
const int M = 1e9 + 7;

//树中两点距离为k的数量
ll dp[50010][510];//以u为根到u距离为j时的方案数量
vector<ll>ed[N];
int n, k;
ll s;

void dfs(int u, int fa)
{
	dp[u][0] = 1;
	for(auto v : ed[u])
	{
		if(v == fa) continue;
		dfs(v, u);
		s += dp[v][k - 1];
		for(int i = 1; i <= k; i ++)
		{
			if(i != k)
		    s += dp[u][i] * dp[v][k - i - 1];
			dp[u][i] += dp[v][i - 1];
		}
	}
}
int main()
{
	cin >> n >> k;
	for(int i = 1; i <= n - 1; i ++)
	{
		int u, v;
		cin >> u >> v;
		ed[u].push_back(v);
		ed[v].push_back(u);
	}
	dfs(1, 0);
	cout << s << endl;
}

PS：这题还可以用点分治，等以后学了在来补上qwq