Tyvj 1052 没有上司的舞会 树形dp

#include<cstdio>
#include<cstring>
#include<vector>
#include<algorithm>
using namespace std;
const int N =6005;
int val[N];
int dp[N][2];
vector<int>g[N];
bool flag[N];

void dfs(int rt){
	if(g[rt].size()==0){
		dp[rt][1]=val[rt];
		dp[rt][0]=0;
		return;
	}
	for(int i=0;i<g[rt].size();i++){
		int u=g[rt][i];
		dfs(u);
		dp[rt][1]+=dp[u][0];
		dp[rt][0]+=max(dp[u][0],dp[u][1]);
	}
	dp[rt][1]+=val[rt];
}
int main(){
	int n;
	while(scanf("%d",&n)!=EOF){
		for(int i=0;i<=n;i++){
			g[i].clear();
		}
		for(int i=1;i<=n;i++){
			scanf("%d",&val[i]);
		}
		int r,c;
		memset(flag,false,sizeof(flag));
		for(int i=1;i<n;i++){
			scanf("%d%d",&r,&c);
			g[c].push_back(r);
			flag[r]=true;
		}
		scanf("%d%d",&r,&c);
		int rt;
		for(int i=1;i<=n;i++){
			if(!flag[i]){
				rt=i;
				break;
			}
		}
		dfs(rt);
		printf("%d\n",max(dp[rt][0],dp[rt][1]));
	}
	return 0;
}

和杭电有一题一样的。