POJ 2342 Anniversary party(树形dp)

题目链接：
POJ 2342 Anniversary party
题意：
有 n 个人要去参加聚会，每个人都不想和他的直接上属出现，给出每个人的从属关系，每个人会有 conviviality ，求聚会的人的 conviviality 和最大是多大？
数据范围： n≤6000
分析：
其实每个人就两种状态 dp[i][0] 不参加， dp[i][1] 参加，如果 i 参加，那么所有 i 的直接下属都不能参加，但是如果 i 不参加，并不意味着他的直接下属必须参加。还要注意可能不止一颗树存在。

#include <stdio.h>
#include <algorithm>
#include <string.h>
#include <math.h>
using namespace std;
typedef long long ll;
const int MAX_N = 6010;

int n, total, root;
int data[MAX_N], head[MAX_N], degree[MAX_N], vis[MAX_N][2], dp[MAX_N][2];

struct Edge {
    int to, next;
}edge[MAX_N];

inline void AddEdge(int from, int to) 
{
    edge[total].to = to;
    edge[total].next = head[from];
    head[from] = total++;
}

int dfs(int x, int flag)
{
    if (vis[x][flag]) return dp[x][flag];
    //printf("x = %d flag = %d\n", x, flag);
    vis[x][flag] = 1;
    int res = 0;
    for (int i = head[x]; i != -1; i = edge[i].next) {
        int to = edge[i].to;
        if (flag == 0) res += max(dfs(to, flag), dfs(to, !flag));
        else res += dfs(to, 0);
    }
    if (flag == 1) res += data[x];
    return dp[x][flag] = res;
}

int main()
{
    while (~scanf("%d", &n)) {
        memset(head, -1, sizeof(head));
        memset(degree, 0, sizeof(degree));
        memset(vis, 0, sizeof(vis));
        total = 0;
        scanf("%d", &data[1]);
        if (n == 0 && data[1] == 0) break;
        for (int i = 2; i <= n; ++i) {
            scanf("%d", &data[i]);
        }
        for (int i = 1; i < n; ++i) {
            int u, v;
            scanf("%d%d", &v, &u);
            AddEdge(u, v);
            degree[v]++;
        }
        int ans = 0;
        for (int i = 1; i <= n; ++i) {
            if (degree[i] == 0) {
                ans += max(dfs(i, 0), dfs(i, 1));
            }
        }
        printf("%d\n", ans);
    }
    return 0;
}