sXSZEp - Online IDE & Debugging Tool

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

#define int long long
const int N = 200005;

int n, w[N];
vector<int> adj[N];
int sub[N];
int best1[N], best2[N]; // 2 nhánh con tốt nhất
int res = 0;

void dfs1(int u, int p) {
    sub[u] = w[u];
    best1[u] = best2[u] = 0;

    for (int v : adj[u]) if (v != p) {
        dfs1(v, u);
        sub[u] += sub[v];
        int val = sub[v];
        if (val > best1[u]) {
            best2[u] = best1[u];
            best1[u] = val;
        } else if (val > best2[u]) {
            best2[u] = val;
        }
    }
}

void dfs2(int u, int p) {
    // cập nhật đáp án tại u
    res = max(res, sub[u] + best1[u] + best2[u]);

    for (int v : adj[u]) if (v != p) {
        // giả sử ta reroot sang v
        int su = sub[u], sv = sub[v];

        // giá trị "ngoài v" khi xem từ u
        int outside = su - sv;
        int candidate = outside;

        // Nếu v chính là best1[u] thì nhánh mạnh nhất ngoài v là best2[u]
        if (sub[v] == best1[u]) candidate = max(candidate, best2[u]);
        else candidate = max(candidate, best1[u]);

        // cập nhật best cho v bằng cách xét nhánh từ u
        int old1 = best1[v], old2 = best2[v];
        if (candidate > best1[v]) {
            best2[v] = best1[v];
            best1[v] = candidate;
        } else if (candidate > best2[v]) {
            best2[v] = candidate;
        }

        // reroot
        sub[v] = sv + (su - sv);

        dfs2(v, u);

        // rollback
        sub[v] = sv;
        best1[v] = old1;
        best2[v] = old2;
    }
}

int32_t main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n;
    for (int i = 1; i <= n; i++) cin >> w[i];
    for (int i = 1; i < n; i++) {
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    dfs1(1, 0);
    dfs2(1, 0);

    cout << res << "\n";
}