代码中第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;
}