Still more natural state.
Order $ f [x] [j] $ $ X $ represents the whole subtree is controlled, and the J $ $ extra layer.
Make $ g [x] [j] $ said they still need $ j $ layer in order to control all $ x $ subtrees.
When transferring $ f [x] [j] $ there are two cases: the control has been previously: $ f [x] [j] + g [y] [j] $, not before control: $ g [x] [j +1] + f [y] [j + 1] $.
Note the shift boundary just fine.
#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; }