洛谷P4657 [CEOI2017]Chase

版权声明：蒟蒻Blog随意转载 https://blog.csdn.net/a1799342217/article/details/81877260

树形DP

洛谷题目传送门

首先对于每个点放磁铁产生的贡献都是固定的。

设$f[i][j],g[i][j]$分别表示从子树$i$中某个点走到$i$和从$i$走到子树中丢了$j$个磁铁的最大值。那么每次转移只要考虑当前点是从哪里来的就好了。

注意更新答案的时候要倒过来再做一遍，顺序是不同的。

代码：

#include<cctype>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define N 100005
#define M 105
#define F inline
using namespace std;
typedef long long LL;
struct edge{ int nxt,to; }ed[N<<1];
int n,m,k,tp,h[N],stk[N];
LL ans,a[N],s[N],f[N][M],g[N][M];
F char readc(){
    static char buf[100000],*l=buf,*r=buf;
    if (l==r) r=(l=buf)+fread(buf,1,100000,stdin);
    return l==r?EOF:*l++;
}
F int _read(){
    int x=0; char ch=readc();
    while (!isdigit(ch)) ch=readc();
    while (isdigit(ch)) x=(x<<3)+(x<<1)+(ch^48),ch=readc();
    return x;
}
#define add(x,y) ed[++k]=(edge){h[x],y},h[x]=k
#define max(x,y) ((x)>(y)?(x):(y))
F void pshp(int x,int v,int fa){
    for (int i=0;i<=m;i++)
        ans=max(ans,f[x][i]+g[v][m-i]);
    for (int i=1;i<=m;i++){
        f[x][i]=max(f[x][i],max(f[v][i],f[v][i-1]+s[x]-a[v]));
        g[x][i]=max(g[x][i],max(g[v][i],g[v][i-1]+s[x]-a[fa]));
    }
}
void dfs(int x,int fa){
    for (int i=1;i<=m;i++)
        f[x][i]=s[x],g[x][i]=s[x]-a[fa];
    for (int i=h[x],v;i;i=ed[i].nxt)
        if ((v=ed[i].to)!=fa) dfs(v,x),pshp(x,v,fa);
    for (int i=h[x],v;i;i=ed[i].nxt)
        if (ed[i].to!=fa) stk[++tp]=ed[i].to;
    for (int i=1;i<=m;i++)
        f[x][i]=s[x],g[x][i]=s[x]-a[fa];
    while (tp) pshp(x,stk[tp--],fa);
}
int main(){
    n=_read(),m=_read();
    for (int i=1;i<=n;i++) a[i]=_read();
    for (int i=1,x,y;i<n;i++){
        x=_read(),y=_read(),add(x,y),add(y,x);
        s[x]+=a[y],s[y]+=a[x];
    }
    return dfs(1,0),printf("%lld\n",ans),0;
}