4557: [JLoi2016]侦察守卫|树形DP

let’s Orz yts大爷

//#pragma comment(linker, "/STACK:20240000,20240000") 
#include<algorithm>
#include<iostream>
#include<cstdlib>
#include<cstring>
#include<cstdio>
#include<vector>
#include<cmath>
#include<ctime>
#include<queue>
#include<set>
#include<map>
#define N 500005
using namespace std;
bool mark[N][22];
int head[N],nxt[N<<1],lst[N<<1];
int f[N][22],g[N][22],w[N];
int n,d,m,tot;
void insert(int x,int y)
{
    lst[++tot]=y; nxt[tot]=head[x]; head[x]=tot; 
    lst[++tot]=x; nxt[tot]=head[y]; head[y]=tot;
}
void dfs(int x,int fa)
{
    int flag=1;
    for(int i=head[x];i;i=nxt[i])
        if(lst[i]!=fa)
        {
            flag=0,dfs(lst[i],x);
            for(int j=1;j<=d;j++)
                mark[x][j]|=mark[lst[i]][j-1];
        }
    if(flag)
    {
        f[x][0]=g[x][0]=w[x]*mark[x][0];
        for(int i=1;i<=d;i++)
            f[x][i]=0,g[x][i]=w[x];
        return ;
    }
    for(int i=1;i<=d;i++)
        for(int j=head[x];j;j=nxt[j])
            if(lst[j]!=fa)
                f[x][i]+=f[lst[j]][i-1];
    for(int i=head[x];i;i=nxt[i])
        if(lst[i]!=fa)
            g[x][d]+=f[lst[i]][d];
    g[x][d]+=w[x];
    for(int i=d-1;i>=0;i--)
    {
        int sum=0;
        for(int j=head[x];j;j=nxt[j])
            if(lst[j]!=fa)
                sum+=f[lst[j]][i];
        g[x][i]=g[x][i+1];
        for(int j=head[x];j;j=nxt[j])
            if(lst[j]!=fa)
                g[x][i]=min(g[x][i],g[lst[j]][i+1]+sum-f[lst[j]][i]);
    }
    f[x][0]=g[x][0];
    for(int i=1;i<=d;i++)
        f[x][i]=min(f[x][i],f[x][i-1]);
    for(int i=d-1;i>=0;i--)
        if(!mark[x][i]) f[x][i]=min(f[x][i],f[x][i+1]);
    g[x][0]=f[x][0];
}

int main()
{
    scanf("%d%d",&n,&d);
    for(int i=1;i<=n;i++)
        scanf("%d",&w[i]);
    scanf("%d",&m);
    for(int i=1;i<=m;i++)
    {
        int x; scanf("%d",&x);
        mark[x][0]=1;
    }
    for(int i=1;i<n;i++)
    {
        int x,y;
        scanf("%d%d",&x,&y);
        insert(x,y);
    }
    dfs(1,0);   
    cout<<f[1][0]<<endl;
    return 0;
}