codeforce 743D. Chloe and pleasant prizes 树dp

D. Chloe and pleasant prizes

题意：一颗以1为根的有根树，每个节点有点权，从中选出2个无相交的子树，使其权值和最大

思路：树dp裸题 dp[u][1] 记录以u为根 选一颗子树的最大值(包括u本身） dp[u][2]记录以u为根 选2颗树的最大值 dp[1][2] 即是答案

AC代码：

#include "iostream"
#include "string.h"
#include "stack"
#include "queue"
#include "string"
#include "vector"
#include "set"
#include "map"
#include "algorithm"
#include "stdio.h"
#include "math.h"
#define ll long long
#define bug(x) cout<<x<<" "<<"UUUUU"<<endl;
#define mem(a) memset(a,0,sizeof(a))
using namespace std;
ll ai[200010];
const ll INF=1e18;
struct Edge{
    int to,next;
};
Edge e[400010];
int head[200010],tot;
void add(ll u, ll v){
    e[tot].to=v;
    e[tot].next=head[u];
    head[u]=tot++;
}
ll sum[200010];
void dfs(ll u, ll fa){
    sum[u]=ai[u];
    for(ll i=head[u]; i!=-1; i=e[i].next){
        ll v=e[i].to;
        if(fa==v) continue;
        dfs(v,u);
        sum[u]+=sum[v];
    }
}
ll dp[200010][5];
void dpans(ll u, ll fa){
    //ll ma=-INF;
    for(ll i=head[u]; i!=-1; i=e[i].next){
        ll v=e[i].to;
        if(v==fa) continue;
        dpans(v,u);
        dp[u][2]=max(dp[u][2],dp[v][2]);
        if(dp[u][1]!=-INF)
            dp[u][2]=max(dp[u][2],dp[v][1]+dp[u][1]);
        dp[u][1]=max(dp[u][1],dp[v][1]);
    }
    dp[u][1]=max(dp[u][1],sum[u]);
}
int main(){
    int n,a,b;
    scanf("%d",&n);
    memset(head,-1,sizeof(head));
    mem(ai),mem(e),mem(sum);
    for(ll i=1; i<=n; ++i){
        scanf("%lld",&ai[i]);
        dp[i][1]=-INF;
        dp[i][2]=-INF;
    }
    for(ll i=1; i<n; ++i){
        scanf("%d%d",&a,&b);
        add(a,b);
        add(b,a);
    }
    dfs(1,1);
    dpans(1,1);     //for(int i=1; i<=n; i++)  cout<<endl<<dp[i][2]<<"  ";   cout<<endl;
    if(dp[1][2]==-INF) {cout<<"Impossible"<<endl;return 0;}
    printf("%lld\n",dp[1][2]);
    return 0;
}