CF1901E 题解

Register_int · 2023-11-29 19:45:58 · 题解

设 dp_u 为 u 子树内的答案，v 为 u 的儿子。先默认所有点都有父亲，那么 u 儿子的删除情况有三种可能：

• 一个都不剩，答案为 a_u。
• 剩了一个，那么 u 度数为 2 会被压缩，答案为 \max dp_v。
• 剩了两个及以上，那么无事发生，选掉前两大后贪心地选非负贡献的儿子即可。

直接把所有儿子的 dp_v 拉下来排序即可做到。接下来考虑如何统计答案：

• 全删光，答案为 0。
• 只剩下一个节点 u，答案为 a_u。
• 根 u 的度数为 1，仅剩一个儿子，答案为 a_u+\max dp_v。
• 根剩下两个儿子，被缩掉了，答案为 dp_v 的前两大之和。
• 根剩下不止两个儿子，强制选前三大后贪心选非负贡献的儿子即可。

时间复杂度 O(n\log n)，瓶颈在于排序。精细实现可做到 O(n)。

AC 代码

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

const int MAXN = 5e5 + 10;

int t, n; ll a[MAXN], dp[MAXN], ans;

vector<int> g[MAXN];

void dfs(int u, int fa) {
    ll x; vector<ll> d; dp[u] = a[u], ans = max(ans, dp[u]);
    for (int v : g[u]) {
        if (v == fa) continue; dfs(v, u), d.push_back(dp[v]);
        ans = max(ans, a[u] + dp[v]), dp[u] = max(dp[u], dp[v]);
    }
    sort(d.begin(), d.end(), greater<ll>());
    if (d.size() >= 2) {
        x = d[0] + d[1], ans = max(ans, x);
        for (int i = 2; i < d.size() && d[i] > 0; i++) x += d[i];
        dp[u] = max(dp[u], a[u] + x);
    }
    if (d.size() >= 3) {
        x = d[0] + d[1] + d[2];
        for (int i = 3; i < d.size() && d[i] > 0; i++) x += d[i];
        ans = max(ans, a[u] + x);
    }
}

int main() {
    for (scanf("%d", &t); t--;) {
        scanf("%d", &n), ans = 0;
        for (int i = 1; i <= n; i++) g[i].clear();
        for (int i = 1; i <= n; i++) scanf("%lld", &a[i]);
        for (int i = 1, u, v; i < n; i++) {
            scanf("%d%d", &u, &v);
            g[u].push_back(v), g[v].push_back(u);
        }
        dfs(1, 0), printf("%lld\n", ans);
    }
}