P3574 [POI2014] FAR-FarmCraft

#include <bits/stdc++.h>
#define IOS ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
#define endl "\n"
#define LL long long
using namespace std;
const int N = 5e5 + 10, M = N * 2;
int n;
int t[N], f[N], g[N];
int h[N], e[M], ne[M], idx;
void add(int a, int b)
{
	e[idx] = b, ne[idx] = h[a], h[a] = idx ++;
}
bool cmp(int a, int b)
{
    return f[a] - g[a] > f[b] - g[b];
}
void dfs(int u, int fa)
{
	f[u] = t[u];
	for (int i = h[u]; i != -1; i = ne[i]) {
		int j = e[i];
		if (j == fa)	continue;
		dfs(j, u);
	}
	vector<int> p;
	for (int i = h[u]; i != -1; i = ne[i]) if (e[i] != fa)  p.push_back(e[i]);
	sort(p.begin(), p.end(), cmp);
	for (auto x : p) {
		f[u] = max(f[u], f[x] + g[u] + 1);
		g[u] += g[x] + 2;
	}
}
int main() {
	IOS;
	cin >> n;
	memset(h, -1, sizeof h);
	for (int i = 1; i <= n; i ++ )	cin >> t[i];
	for (int i = 1; i < n; i ++ ) {
		int a, b;
		cin >> a >> b;
		add(a, b), add(b, a);
	}
	dfs(1, -1);
	cout << max(f[1], g[1] + t[1]);
	return 0;
}