1901E - Compressed Tree - CodeForces Solution

#include <bits/stdc++.h>

using namespace std;

void solve() {
    int n;
    cin >> n;
    
    vector <int> a(n + 1);
    for (int i = 1; i <= n; i++) {
        cin >> a[i];
    }
    
    vector < vector <int> > g(n + 1);
    for (int i = 1; i < n; i++) {
        int u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    
    const long long inf = 1e18;
    
    vector <vector <long long> > dp(n + 1, vector <long long>(4, -inf));
    
    function <void(int, int)> dfs = [&](int v, int p) {
        dp[v][0] = a[v];
        for (int u : g[v]) {
            if (u == p) continue;
            dfs(u, v);
            vector <long long> cur(4, -inf);
            for (int c1 = 0; c1 <= 3; c1++) {
                for (int c2 = 0; c2 <= 3; c2++) {
                    int c = min(c1 + 1, 3);
                    long long res = dp[v][c1] + dp[u][c2];
                    if (c1 == 2) {
                        res += a[v];
                    }
                    if (c1 == 1) {
                        res -= a[v];
                    }
                    if (c2 == 2) {
                        res += a[u];
                    }
                    if (c2 == 1) {
                        res -= a[u];
                    }
                    cur[c] = max(cur[c], res);
                }
            }
            for (int i = 0; i < 4; i++) {
                dp[v][i] = max(dp[v][i], cur[i]);
            }
        }
    };
    
    dfs(1, -1);
    
    long long ans = 0;
    for (int i = 1; i <= n; i++) {
        for (int j = 0; j < 4; j++) {
            ans = max(ans, dp[i][j]);
        }
    }
    
    cout << ans << endl;
}

int main() {
    ios::sync_with_stdio(0);
    int tt = 1;
    cin >> tt;
    while (tt--) solve();
}

/*
 1
 4
 1 -2 2 1
 1 2
 3 2
 2 4
 */