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 105
int 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;
}