POI Task FarmCraft

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

vector<int> c;
vector<vector<int>> g;

// returns pair (t, c)
// t = time spent in subtree, c = time after entering subtree that all nodes are installed
pair<int, long long> dfs(int node, int parent = -1) {
    vector<pair<int, long long>> v;
    for (int child : g[node]) {
        if (child == parent) continue;
        pair<int, long long> pr = dfs(child, node);
        pr.first += 2, pr.second += 1;
        v.push_back(pr);
    }
    sort(v.begin(), v.end(), [](const pair<int, long long> &a, const pair<int, long long> &b) {
        return a.first-a.second < b.first-b.second; // sort by t-c
    });
    int sumT = 0;
    long long maxSetup = 0;
    for (size_t i = 0; i < v.size(); sumT += v[i++].first) {
        maxSetup = max(maxSetup, sumT + v[i].second);
    }
    maxSetup = max(maxSetup, (long long)c[node] + (node ? 0 : sumT)); // node's setup time
    return pair<int, long long>{sumT, maxSetup};
}

int main() {
    ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    int n; cin >> n;
    c.resize(n);
    for (int i = 0; i < n; ++i) cin >> c[i];
    g.resize(n);
    for (int i = 1; i < n; ++i) {
        int u, v; cin >> u >> v;
        --u, --v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    cout << dfs(0).second << endl;
    return 0;
}