【XSY2411】【CF161D】Distance in Tree（树形dp）

(Description)

一棵树是一个不含环的连通图。

树上两点间的距离是两点间最短路径的长度（边的）。

给你一棵(n)个点的树和一个正整数(k)。找出不同的距离为(k)的点对的数量。注意点对(（u,v）)和(（v,u）)被认为是相同的。

---

(Input)

第一行包含两个整数(n)和(k) (（1≤n≤50000,1≤k≤500）)——点数和要求的两点间距离。

接下来(n-1)行代表了边(“a_{i}) (b_{i}”)（没有引号）(（1≤a_{i},b_{i}≤n,a_{i}≠b_{i}）)，(a_{i})和(b_{i})是第(i)条边的端点。所有给出的边都是不同的。

---

(Output)

输出单独一个整数——不同的距离正好为(k)的树上点对的数量。

---

(Sample Input)

样例输入1
5 2
1 2
2 3
3 4
2 5

样例输入2
5 3
1 2
2 3
3 4
4 5

---

(Sample Output)

样例输出1
4

样例输出2
2

---

(HINT)

第一个样例的距离为(2)的点对是(（1,3）)，(（1,5）)，(（3,5）)和(（2,4）)。

---

(Source)

 练习题　树1-树形DP

---

思路

我们考虑计算出对于以(u)为根的子树中，距离(u)各种距离的节点的个数，这个值可以通过(u)的儿子(v)来更新

所以我们选择树形(dp)

我们设一个数组(sum[u][t])，表示以(u)为根的子树中，距离(u)的距离为(t)的节点的个数

那么答案(ans)就是(sum_{j=0}^{k-1}sum[u][j] imes sum[v][k-j-1])

这个简单解释下：枚举与(u)距离(j)的点，因为要统计距离为(k)的点个数，于是我们还需要距离(k-j)的点，因为通过(v)转移 ，有一个距离(1)，所以只需要找与(v)距离(k-j-1)的点的个数，乘起来就可以了

那么又因为在计算(v)的时候，还没有将(v)的(sum)统计进(u)的(sum)中，所以在计算(ans)的时候不会有重复统计

接下来，来讲讲怎么转移方程

也十分简单：(sum_{j=1}^{k}sum[u][j]+=sum[v][j-1])

很好理解那就不细讲了

于是我们就可以一遍(dfs)处理出(sum)数组和答案，最后直接输出就好了

---

代码

#include<bits/stdc++.h>
using namespace std;
const int N=50010;
int n,k,cnt=0;
int to[N<<1],nxt[N<<1],head[N];
int sum[N][510];
int ans=0;
inline void add(int u,int v)
{
	to[++cnt]=v;
	nxt[cnt]=head[u];
	head[u]=cnt;
}
void dfs(int u,int fa)
{
	sum[u][0]=1;
	for(int i=head[u];i;i=nxt[i])
	{
		int v=to[i];
		if(v==fa)continue;
		dfs(v,u);
		for(int j=0;j<k;j++)ans+=sum[u][j]*sum[v][k-j-1];
		for(int j=1;j<=k;j++)sum[u][j]+=sum[v][j-1];
	}
}
inline int read()
{
	int x=0,f=1;
	char ch=getchar();
	while(!isdigit(ch))
	{
		if(ch=='-')f=-1;
		ch=getchar();
	}
	while(isdigit(ch))
	{
		x=(x<<3)+(x<<1)+(ch^48);
		ch=getchar();
	}
	return x*f;
}
int main()
{
	n=read(),k=read();
	int a,b;
	for(int i=1;i<n;i++)
	{
		a=read(),b=read();
		add(a,b);add(b,a);
	}
	dfs(1,0);
	printf("%d",ans);
	return 0;
}