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