Codeforces 161D 树形dp

stdforces

于 2022-07-10 23:18:24 发布

阅读量201

点赞数

文章标签：算法

本文链接：https://blog.csdn.net/stdforces/article/details/125712635

版权

给一棵树，问树上距离为k的点对有多少个？(u,v)与(v,u)这样的算一种

一开始思考很简单，自然就想到dp，设 $f[u][k]$ 为 $u$ 结点向下 $k$ 个位置有多少个结点

转移也很容易想出来:

遍历所有向下能到达的节点v， $f[u][k]+=f[v][k-1]$

但这只适用与 $u$ 只有一棵子树，因为这样的话两棵子树间无法计数

画了个有两个子树的样例，加多一个数组 $g$ ，含义与前面一致，只不过两个分别记录两棵子树。

然后又发现如果一个结点有n棵子树，那么这个方法就不适用了。

正解是，在dfs内，每次回来更新数组必然是完成了一棵子树的统计，而此时 $f$ 已经保存了之前所有子树的答案，这个时候先拿 $v$ 的 $f$ 与 $u$ 的 $f$ 计算答案，再用 $v$ 的 $f$ 更新 $u$ 的 $f$

统计代码

void dfs(int u,int fa)
{
	f[u][0]=1;
	for(int i=head[u];i;i=edge[i].next)
	{
		int v=edge[i].to;
		if(v==fa) continue;
		dfs(v,u);
		for(int j=0;j<=k;j++) ans+=f[u][j]*f[v][k-j-1];
		for(int j=1;j<=k;j++) f[u][j]+=f[v][j-1];
	}
}

比如现在刚统计完 $v$ ，我们刚回溯还没更新 $u$ 的 $f$ 时， $u$ 的 $f$ 只保存了 $v_{1},v_{2}$ 两棵子树的点，我们就拿 $v_{3}$ 的 $f$ 与前两棵子树计算，前两棵子树往下 $i$ 的点数与 $v_{3}$ 往下 $j$ 的点数就是 $v_{3}$ 为出发点到 $v_{1},v_{2}$ 的点距离为 $i+j$ 的点对，因为每一棵子树都是与前面的子树做计算，所以次数不会计算重复。

#include<bits/stdc++.h>
#define ll long long
using namespace std;

ll read()
{
	ll ret=0,base=1;
	char ch=getchar();
	while(!isdigit(ch))
	{
		if(ch=='-') base=-1;
		ch=getchar();
	}
	while(isdigit(ch))
	{
		ret=(ret<<3)+(ret<<1)+ch-48;
		ch=getchar();
	}
	return ret*base;
}

struct way
{
	int to,next;
}edge[200005];
int cnt,head[100005];

void add(int u,int v)
{
	edge[++cnt].to=v;
	edge[cnt].next=head[u];
	head[u]=cnt;
}

void add_(int u,int v){add(u,v);add(v,u);}

int n,k,f[100005][505],ans;

void dfs(int u,int fa)
{
	f[u][0]=1;
	for(int i=head[u];i;i=edge[i].next)
	{
		int v=edge[i].to;
		if(v==fa) continue;
		dfs(v,u);
		for(int j=0;j<=k;j++) ans+=f[u][j]*f[v][k-j-1];
		for(int j=1;j<=k;j++) f[u][j]+=f[v][j-1];
	}
}

int main()
{
	n=read();k=read();
	for(int i=1;i<n;i++) add_(read(),read());
	dfs(1,0);
	cout<<ans;
	return 0;
}