AcWing 285. 没有上司的舞会

AcWing 285. 没有上司的舞会

树形dp

从小偷系列问题演变过来的树形dp问题

#include<iostream>
#include<cstring>
using namespace std;
const int N=10010;

int n;
int h[N],e[N],ne[N],happy[N],idx;
int f[N][2];
bool has_fa[N];


void add(int a,int b){
    e[idx]=b;
    ne[idx]=h[a];
    h[a]=idx++;
}

void dfs(int u){
    f[u][1]=happy[u];
    
    for(int i=h[u];i!=-1;i=ne[i]){
        int t=e[i];
        dfs(t);
        f[u][0]+=max(f[t][0],f[t][1]); //not select cur
        f[u][1]+=f[t][0];  //select cur
    }
    
}

int main(){
    memset(h,-1,sizeof h);
    scanf("%d",&n);
    
    for(int i=1;i<=n;i++)scanf("%d",&happy[i]);
    
    for(int i=1;i<n;i++){
        int a,b;
        scanf("%d%d",&a,&b);
        add(b,a);
        
        has_fa[a]=true;
    }
    
    int root=1;
    while(has_fa[root])root++;
    
    dfs(root);
    
    
    printf("%d",max(f[root][0],f[root][1]));
    
    
}