没有上司的舞会

#include<iostream>
#include<cstring>
#include<algorithm>

using namespace std;

const int N = 6010;

int hp[N];
int e[N],h[N],ne[N],idx;
int f[N][2];
bool hf[N];
int n;

//邻接表建树
void add(int a,int b){
    e[idx] = b;
    ne[idx] = h[a];
    h[a] = idx ++;
}

void dfs(int u){
    //dp赋值
    f[u][1] = hp[u];

    //遍历子节点
    for(int i = h[u];i != -1;i = ne[i]){
        int j = e[i];
        //子节点更新
        dfs(j);

        //选不选父节点
        f[u][1] += f[j][0];
        f[u][0] += max(f[j][1],f[j][0]);
    }
}

int main(){
    cin >> n;

    for(int i = 1;i <= n;i ++) cin >> hp[i];
    memset(h,-1,sizeof h);

    for(int i = 1;i < n;i ++){
        int a,b;
        cin >> a >> b;

        //b为a的父亲
        add(b,a);
        hf[a] = 1;
    }

    //找根节点
    int rt = 1;
    while(hf[rt]) rt ++;

    //根节点开搜
    dfs(rt);

    cout << max(f[rt][0],f[rt][1]);
}