POJ 2486 Apple Tree ( 树型DP )

#include <iostream>
#include <cstring>
#include <deque>
using namespace std;

#define SIZE 230
#define BACK 1
#define AWAY 0

int DP[SIZE][SIZE][2];
bool visits[SIZE];
int vals[SIZE];
deque< int > tree[SIZE];

int num, steps;

void dfs( int u ){

    visits[u] = true;
    const int len = tree[u].size();

    for( int i = 0; i <= steps; ++i )
        DP[u][i][BACK] = DP[u][i][AWAY] = vals[u];

    for( int i = 0; i < len; ++i ){

        int son = tree[u][i];

        if( visits[son] )
            continue;

        dfs( son );

        for( int s = steps; s >= 0; --s ){
            for( int ss = 0; ss <= s; ++ss ){
                /*
                    从 u 出发，回到 u，需要多走两步 u->son,son->u,
                    分配给 son 子树 ss 步，其他子树 s - ss 步，都返回.
                */
                DP[u][s + 2][BACK] = max( DP[u][s + 2][BACK],
                                          DP[u][s - ss][BACK] + DP[son][ss][BACK] );

                /*
                    不回 u （去 u 的其他子树），在 son 返回.
                */
                DP[u][s + 2][AWAY] = max( DP[u][s + 2][AWAY],
                                          DP[u][s - ss][AWAY] + DP[son][ss][BACK] );

                /*
                    先遍历 u 的其他子树，回到 u 后，遍历 son 子树，
                    在当前子树 son 不返回，多走一步.
                */
                DP[u][s + 1][AWAY] = max( DP[u][s + 1][AWAY],
                                          DP[u][s - ss][BACK] + DP[son][ss][AWAY] );

            }
        }
    }
}

int main(){

    int u, v;

    while( cin >> num >> steps ){

        memset( DP,     0,     sizeof( DP ) );
        memset( visits, false, sizeof( visits ) );

        for( int i = 1; i <= num; ++i )
            tree[i].clear();

        for( int i = 1; i <= num; ++i )
            cin >> vals[i];

        for( int i = 1; i <= num - 1; ++i ){

            cin >> u >> v;

            tree[u].push_back( v );
            tree[v].push_back( u );

        }

        dfs( 1 );

        cout << max( DP[1][steps][BACK], DP[1][steps][AWAY] ) << endl;

    }

    return 0;

}