hdu 1520 Anniversary party 树形DP

#include<stdio.h>
#include<string.h>
#include<vector>
#include<algorithm>
using namespace std;

int n,dp[6000+5][2],bin[6000+5];
vector<int>d[6000+5];

void tree_dp(int x)
{
    int i,len=d[x].size();
    for(i=0;i<len;i++)
    {
        tree_dp(d[x][i]);
        dp[x][0]+=max(dp[d[x][i]][0],dp[d[x][i]][1]);
        dp[x][1]+=dp[d[x][i]][0];
    }
    return ;
}

int main()
{
    int i,a,b,rt;
    while(~scanf("%d",&n))
    {
        memset(dp,0,sizeof(dp));
        memset(bin,0,sizeof(bin));
        for(i=0;i<=6000;i++)
            d[i].clear();
        for(i=1;i<=n;i++)
            scanf("%d",&dp[i][1]);
        while(~scanf("%d%d",&a,&b))
        {
            if(a==0&&b==0) break;
            rt=b;
            bin[a]=b;
            d[b].push_back(a);
        }
        while(bin[rt]) rt=bin[rt];
        tree_dp(rt);
        printf("%d\n",max(dp[rt][0],dp[rt][1]));
    }
    return 0;
}