poj 2486 Apple Tree(树形dp)

/*
    说一下做题感受，首先一看就可以确定是一道树形背包问题。
    刚开始想的是两个状态来表示(rt,i),但这样的话不好表示从一个子节点回来再去访问另外子结点的情况，
    所以就想到了三个状态(rt,i,0)和(rt,i,1),分别表示是否回到rt结点的最大值，这样就可以写出状态方程：
    (rt,i,0)=max{(rt,i,0),(s,j,0)+(rt,i-2-j,0)};  i>=2;i-2>=j;
    (rt,i,1)=max{(rt,i,1),(s,j,0)+(rt,i-2-j,1),(s,j,1)+(rt,i-1-j,0)};
*/

#include<cstdio>#include<cstring>#include<cstdlib>#include<iostream>#include<vector>using namespace std;#define maxn 105int dp[maxn][2*maxn][2],vis[maxn],k,val[maxn];vector<int>sons[maxn];void dfs(int rt){    vis[rt]=1;    for(int i=0;i<=k;i++) dp[rt][i][0]=dp[rt][i][1]=val[rt];    for(int u=0;u<sons[rt].size();u++){        int s=sons[rt][u];        if(vis[s]) continue;        dfs(s);        for(int i=k;i>=1;i--)            for(int j=0;j<i;j++){                //先更新(rt,i,1)再更新(rt,i,0)                dp[rt][i][1]=max(dp[rt][i][1],dp[s][j][1]+dp[rt][i-1-j][0]);                if(i>=2+j) dp[rt][i][1]=max(dp[rt][i][1],dp[s][j][0]+dp[rt][i-2-j][1]); //刚开始少了这步,会少了很多情况,因为需要兼顾前面的情况                if(i>=2+j) dp[rt][i][0]=max(dp[rt][i][0],dp[s][j][0]+dp[rt][i-2-j][0]);            }    }}void init(){    for(int i=0;i<maxn;i++) sons[i].clear();    memset(vis,0,sizeof(vis));    memset(dp,0,sizeof(dp));}void print(){    int ans=0;    for(int i=0;i<=k;i++) ans=max(ans,dp[1][i][1]),ans=max(ans,dp[1][i][0]);    cout<<ans<<endl;}int main(){    int i,j,n,a,b;    while(cin>>n>>k){        init();        for(i=1;i<=n;i++) cin>>val[i];        for(i=1;i<n;i++){            cin>>a>>b;            sons[a].push_back(b);            sons[b].push_back(a);        }        dfs(1);        print();    }    return 0;}