POJ 2486 Apple Tree 树形DP

题意：一个小女孩到了一颗苹果树上。苹果树上每个节点都有不同数量的苹果。现在小女孩从根结点出发，最多走K步，问最多可以吃到多少苹果？
思路：树形DP，和前面一道题同样的思路。
定义dp[u][j][1]为从节点u出发，在以u为根的子树中行走了j步，最终留在了子树中。
dp[u][j][0]为从节点u出发，在以u为根的子树中行走了j步，最终没有留在子树中。
下面考虑状态转移：
1.对于dp[u][j][0],只有一种情况

$$dp[u][j][0] = max(dp[v][k][0] + dp[u][j-k-2][0]), v \in son(u)，0 \le k \le j-2$$
2.对于dp[u][j][1],对于u的儿子v,我们可以选择留在v中，也可以选择不留在其中
$$dp[u][j][1] = max(dp[v][k][0] + dp[u][j-k-2][1], dp[v][k][1] + dp[u][j-k-1][0])$$
边界条件：
$$dp[u][0][0] = dp[u][0][1] = apple_u$$
最终的答案是 $max(dp[1][j][1]), 0 \le j \le K$
代码如下：

#include <cstdio>
#include <cstring>
#include <algorithm>

using namespace std;

const int MAX = 210;

int N,S,K;
int to[MAX<<1],nxt[MAX<<1],head[MAX],tot;
int dp[MAX][MAX][2],ap[MAX];


void init()
{
    memset(head,-1,sizeof(head));
    memset(dp,0,sizeof(dp));
    tot = 0;
}

void addedge(int u, int v)
{
    to[tot] = v;
    nxt[tot] = head[u];
    head[u] = tot++;
}

void dfs(int u, int p)
{
    dp[u][0][0] = dp[u][0][1] = ap[u];
    for(int i = head[u]; ~i; i = nxt[i]){
        int v = to[i];
        if(v == p) continue;
        dfs(v,u);
        for(int j = K; j >= 1; --j){
            for(int k = 0; k <= j - 1; ++k){
                if(j-k-2 >= 0){
                    dp[u][j][0] = max(dp[u][j][0],dp[v][k][0] + dp[u][j-k-2][0]);
                    dp[u][j][1] = max(dp[u][j][1],dp[v][k][0] + dp[u][j-k-2][1]);
                }
                dp[u][j][1] = max(dp[u][j][1],dp[v][k][1] + dp[u][j-k-1][0]);
            }
        }
    }
}

int main(void)
{
    //freopen("input.txt","r",stdin);
    while(~scanf("%d%d",&N,&K)){
        init();
        for(int i = 1; i <= N; ++i)
            scanf("%d",&ap[i]);
        for(int i = 1; i <= N - 1; ++i){
            int u, v;
            scanf("%d%d",&u,&v);
            addedge(u,v);
            addedge(v,u);
        }
        dfs(1,0);
        int ans = 0;
        for(int i = 0; i <= K; ++i)
            ans = max(ans,dp[1][i][1]);
        printf("%d\n",ans);
    }
    return 0;
}