HDU6540 Neko and tree

http://acm.hdu.edu.cn/showproblem.php?pid=6540

Problem Description
Neko has a tree with n nodes.
There are m key nodes on tree. Neko want to you to selecte some key nodes satisfying the distance between any two selected nodes less than or equal to k.
Neko thinks this work is too easy,so Neko want to know how many different way for selecting nodes.
Calculate the answer after mod 109+7.
Note that you have to select at least one key node.

题意：一颗树，n个结点，其中m个关键结点，要求选 $只包含关键结点的结点集，其中任意两点间距离都不超过k$只包含关键结点的结点集，其中任意两点间距离都不超过k，求选的方案数。
思路：树形dp，设 $f(u,i):以u为根的子树，在满足任意两点间距离不超过k的前提下，结点集中距离u最远的点距离u为i的方案数$f(u,i):以u为根的子树，在满足任意两点间距离不超过k的前提下，结点集中距离u最远的点距离u为i的方案数
用泛化背包的思想更新。
需要用前缀和优化复杂度。读了这篇https://blog.csdn.net/cat_pikapikapi/article/details/90521780

#include<bits/stdc++.h>
using namespace std;
#define maxn 5010
#define mod 1000000007
#define ll long long
#define M(x) ((x)%=mod)

int n,m,k;
ll f[maxn][maxn],pre[maxn][maxn],g[maxn],is[maxn];
ll ans;
vector<int> G[maxn];

void cal(int u,int v) //将子树v加入u的更新操作
{										
	for(int i=1;i<k;i++)g[i]=f[u][i]*pre[v][min(i,k-i)-1],M(g[i]);
	for(int i=1;i<k;i++)g[i]+=pre[u][min(i,k-i)]*f[v][i-1],M(g[i]);
	for(int i=1;i<=k/2;i++)g[i]-=f[u][i]*f[v][i-1],g[i]=(g[i]+mod)%mod;//最深在u+最深在v-uv同样深
	for(int i=1;i<=k;i++)
	{		//只选择v+uv搭配
		f[u][i]+=f[v][i-1]+g[i];M(f[u][i]);
		pre[u][i]=pre[u][i-1]+f[u][i]; M(pre[u][i]);
	}
}

void dfs(int u,int fa)
{
	for(int i=0;i<G[u].size();i++)
	{
		int v=G[u][i];
		if(v==fa)continue;
	    dfs(v,u);
		cal(u,v);
	}
	if(is[u]){	//u本身是关键点，那么子树内方案树加倍，因为加一个u结点仍满足条件
		f[u][0]=pre[u][0]=1;
		for(int i=1;i<=k;i++)
		{
			f[u][i]*=2;M(f[u][i]);
			pre[u][i]=pre[u][i-1]+f[u][i];M(pre[u][i]);
		}
	}
	if(u==1)ans+=pre[u][k];	//非根结点现在只算恰好为k的，别的后续更新。
	else ans+=f[u][k];
	M(ans);
}

int main()
{
	//freopen("input.in","r",stdin);
	int x,y;
	cin>>n>>m>>k;
	for(int i=1;i<n;i++)
	{
		scanf("%d%d",&x,&y);
		G[x].push_back(y);
		G[y].push_back(x);
	}
	for(int i=1;i<=m;i++)cin>>x,is[x]=1;
	dfs(1,0);
	cout<<ans<<endl;
	return 0;
}