【洛谷 P1352】没有上司的舞会

树形dp

#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
const int N=6010;
int f[N][2],w[N],in[N],out[N],n;
int head[N],nxt[N],son[N],size;
void uni(int x,int y){
    size++;
    nxt[size]=head[x];
    head[x]=size;
    son[size]=y;
    in[x]++;
    out[y]++;
}
int dp(int x,int k){
    int &ans=f[x][k];
    if (ans)
        return ans;
    if (!in[x]){
        if (!k)
            return ans=0;
        else
            return ans=w[x];
    }
    if (k)
        ans=w[x];
    else
        ans=0;
    if (k)
        for (int e=head[x];e;e=nxt[e]){
            int y=son[e];
            ans+=dp(y,1-k);
        }
    else
        for (int e=head[x];e;e=nxt[e]){
            int y=son[e];
            int r1=dp(y,1-k);
            int r2=dp(y,k);
            ans+=max(r1,r2);
        }
    return ans;
}
int main(){
    int k,l;
    memset(f,0,sizeof(f));
    size=0;
    scanf("%d",&n);
    for (int i=1;i<=n;i++)
        scanf("%d",&w[i]);
    for (int i=1;i<=n;i++){
        scanf("%d %d",&l,&k);
        if (l&&k)
            uni(k,l);
    }
    for (int i=1;i<=n;i++)
        if (!out[i])
            uni(0,i);
    printf("%d",dp(0,0));
    return 0;
}