KWqf2m - Online C++0x Compiler & Debugging Tool

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

const int N = 1000005;
int n;
ll w[N], sub[N];
vector<int> g[N];
ll ans = 0;

void dfs1(int u, int p) {
    sub[u] = w[u];
    for (int v : g[u]) if (v != p) {
        dfs1(v,u);
        sub[u] += sub[v];
    }
}

void dfs2(int u, int p) {
    ll best1=0, best2=0; int id1=-1;
    for (int v : g[u]) if (v != p) {
        if (sub[v] > best1) {best2=best1; best1=sub[v]; id1=v;}
        else if (sub[v] > best2) best2=sub[v];
    }
    ans = max(ans, sub[u]);
    for (int v : g[u]) if (v != p) {
        ll oldu=sub[u], oldv=sub[v];
        if (v==id1) {
            sub[u]=w[u]+best2;
        }
        sub[v]+=sub[u];
        dfs2(v,u);
        sub[u]=oldu; sub[v]=oldv;
    }
}

int 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;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs1(1,-1);
    dfs2(1,-1);
    cout<<ans<<"\n";
}