2021 icpc 南京 h 树形dp

代码中第55行那个位置我没想到，导致wa了几发。
这个题被卡在了t=3的讨论上面，其实把式子推出来发现只和最大值有关系，最大值必选。
子节点的选择和子节点为头节点的子树的选择无关。

#include <bits/stdc++.h>
using namespace std;
#define ll long long

const int N = 1e5+10, M = 2*N;

int n;
int h[N], ne[M], e[M], idx;
ll a[N], t[N];
ll sum[N], f[N];

void init() {
    idx = 0;
    for(int i = 1;i <= n;i ++) {
        h[i] = -1;
        sum[i] = 0;
        f[i] = 0;
    }
}

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

void dfs(int u, int fa) {
    for(int i = h[u];i != -1;i = ne[i]) {
        int j = e[i];
        if(j != fa) dfs(j, u);
    }
    for(int i = h[u];i != -1;i = ne[i]) {
        int j = e[i];
        if(j == fa) continue;
        sum[u] += f[j];
    }

    ll maxl = 0, id;
    for(int i = h[u];i != -1;i = ne[i]) {//第一种情况 只选择一个
        int j = e[i];
        if(j == fa) continue;
        if(t[j] == 3 && a[j] > maxl) {
            maxl = a[j];
            id = j;
        }
        f[u] = max(f[u], sum[u] + a[j]);
    }

    if(maxl == 0) return ;

    for(int i = h[u];i != -1;i = ne[i]) {//任意选择一个，然后选择t为3的最大值
        int j = e[i];
        if(j == fa || j == id) continue;
        f[u] = max(f[u], sum[u] - f[j] + a[j] + sum[j] + a[id]);
    }    
    for(int i = h[u];i != -1;i = ne[i]) {//先选择t为3的最大值，然后再选一个t为3的值
        int j = e[i];
        if(j == fa || j == id || t[j] != 3) continue;
        f[u] = max(f[u], sum[u] - f[id] + a[id] + sum[id] + a[j]);
    }
    return ;
}

int main() {
    int Case, l, r;
    scanf("%d", &Case);
    while(Case --) {
        scanf("%d", &n);
        init();
        for(int i = 1;i <= n;i ++) scanf("%lld", &a[i]);
        for(int i = 1;i <= n;i ++) scanf("%lld", &t[i]);
        for(int i = 1;i < n;i ++) {
            scanf("%d%d", &l,&r);
            add(l, r); add(r, l);
        }
        dfs(1, 0);
        printf("%lld\n", f[1] + a[1]);
    }
    return 0;
}