状态还比较自然.
令 $f[x][j]$ 表示 $x$ 的子树全被控制,且多出来 $j$ 层.
令 $g[x][j]$ 表示还需要 $j$ 层才能控制 $x$ 所有子树.
转移 $f[x][j]$ 的时候有两种情况:之前已经控制:$f[x][j]+g[y][j]$,之前没控制:$g[x][j+1]+f[y][j+1]$.
注意一下转移边界就好了.
#include <cstdio>
#include <string>
#include <cstring>
#include <algorithm>
#define N 500002
#define inf 1000000000
using namespace std;
void setIO(string s)
{
string in=s+".in";
string out=s+".out";
freopen(in.c_str(),"r",stdin);
// freopen(out.c_str(),"w",stdout);
}
int D,n,edges;
int f[N][21],g[N][21],v[N],vis[N],hd[N],to[N<<1],nex[N<<1];
void add(int u,int v)
{
nex[++edges]=hd[u],hd[u]=edges,to[edges]=v;
}
void dfs(int x,int ff)
{
if(vis[x]) f[x][0]=g[x][0]=v[x];
int i,j;
for(i=1;i<=D;++i) f[x][i]=v[x];
f[x][D+1]=inf;
for(i=hd[x];i;i=nex[i]) if(to[i]!=ff)
{
int y=to[i];
dfs(y,x);
for(j=0;j<=D;++j)
f[x][j]=min(f[x][j]+g[y][j],f[y][j+1]+g[x][j+1]);
for(j=D;j>=0;--j) f[x][j]=min(f[x][j],f[x][j+1]);
g[x][0]=f[x][0];
for(j=1;j<=D;++j) g[x][j]+=g[y][j-1];
for(j=1;j<=D;++j) g[x][j]=min(g[x][j],g[x][j-1]);
}
}
int main()
{
// setIO("input");
int i,j,m;
scanf("%d%d",&n,&D);
for(i=1;i<=n;++i) scanf("%d",&v[i]);
scanf("%d",&m);
for(i=1;i<=m;++i)
{
int x;
scanf("%d",&x),vis[x]=1;
}
for(i=1;i<n;++i)
{
int x,y;
scanf("%d%d",&x,&y),add(x,y),add(y,x);
}
dfs(1,0);
printf("%d\n",g[1][0]);
return 0;
}