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

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

using ll = long long;

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

    int n;
    if(!(cin >> n)) return 0;
    vector<ll> w(n+1);
    for (int i = 1; i <= n; ++i) cin >> w[i];
    vector<vector<int>> g(n+1);
    g.reserve(n+1);
    for (int i = 0; i < n-1; ++i) {
        int u, v; cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }

    auto bfs_far = [&](int s, vector<int>* parent_out=nullptr)->int{
        vector<int> par(n+1, -1);
        queue<int> q; q.push(s); par[s]=0;
        int last=s;
        while(!q.empty()){
            int u=q.front(); q.pop();
            last=u;
            for(int v: g[u]) if(par[v]==-1){
                par[v]=u; q.push(v);
            }
        }
        if(parent_out) *parent_out = move(par);
        return last;
    };

    vector<int> par1;
    int a = bfs_far(1);
    int b = bfs_far(a, &par1);

    // reconstruct a-b path
    vector<int> path;
    for(int x=b; x!=0; x=par1[x]) path.push_back(x);
    // path now from b -> ... -> a ; reverse to a..b
    reverse(path.begin(), path.end());
    int dlen = (int)path.size()-1; // số cạnh của đường kính

    // ---- helpers for component DP (không đệ quy) ----
    struct CompRes {
        int height;   // độ sâu lớn nhất từ root đến lá trong component
        ll sumBest;   // Sum = tổng G dọc đường tốt nhất
        ll valBest;   // Val = tổng k*G (k bắt đầu từ 1) dọc đường tốt nhất
    };

    auto solve_component = [&](int root, int forbidParent)->CompRes{
        // iterative postorder on the component that does NOT cross (root <-> forbidParent)
        vector<int> order;
        order.reserve(1024);
        vector<int> parent(n+1, -2); // -2 = unseen; -1 = no parent (root)
        vector<int> stk; stk.reserve(1024);
        // stack of (u, iterator index, entered?)
        vector<int> itIdx(n+1, 0);

        // manual stack for postorder
        parent[root] = -1;
        stk.push_back(root);
        vector<int> visited;
        visited.reserve(1024);
        while(!stk.empty()){
            int u = stk.back();
            if(itIdx[u] == 0) visited.push_back(u); // first time we see u
            if(itIdx[u] < (int)g[u].size()){
                int v = g[u][itIdx[u]++];
                if(v == parent[u]) continue;
                if(u==root && v==forbidParent) continue; // cut the center edge
                parent[v] = u;
                stk.push_back(v);
            }else{
                order.push_back(u);
                stk.pop_back();
            }
        }
        // compute subsum (=G), height, DP
        unordered_set<int> inComp(visited.begin(), visited.end());
        // arrays only for nodes in component
        // to avoid allocating n-sized arrays each call, use unordered_map-like via vectors + check set
        // but for speed and memory, we use arrays of size n+1 reused; mark only used nodes.
        static vector<ll> subsum, sumBest, valBest;
        static vector<int> height;
        if((int)subsum.size()!=n+1){
            subsum.assign(n+1,0);
            sumBest.assign(n+1,0);
            valBest.assign(n+1,0);
            height.assign(n+1,0);
        }
        for(int u: order){
            ll s = w[u];
            int h = 0;
            for(int v: g[u]){
                if(v==parent[u]) continue;
                s += subsum[v];
                h = max(h, height[v]+1);
            }
            subsum[u] = s;               // G(u) trong component
            height[u] = h;

            if(h==0){
                sumBest[u] = subsum[u];
                valBest[u] = subsum[u];  // k=1
            }else{
                // chọn child trên nhánh sâu nhất cho giá trị lớn nhất
                ll bestValShift = LLONG_MIN;
                ll ch_sum = 0, ch_val = 0;
                for(int v: g[u]){
                    if(v==parent[u]) continue;
                    if(height[v]+1 == h){
                        ll cand = valBest[v] + sumBest[v];
                        if(cand > bestValShift){
                            bestValShift = cand;
                            ch_sum = sumBest[v];
                            ch_val = valBest[v];
                        }
                    }
                }
                sumBest[u] = subsum[u] + ch_sum;
                valBest[u] = subsum[u] + bestValShift; // 1*G(u) + (val_child + sum_child)
            }
        }
        if(inComp.find(root)==inComp.end()){
            // component is empty (shouldn't happen)
            return {0, 0, 0};
        }
        return {height[root], sumBest[root], valBest[root]};
    };

    auto solve_center_vertex = [&](int c)->long long{
        // Root whole tree at c to get subsum (=G) and DP for all nodes
        // iterative postorder from c
        vector<int> parent(n+1, -2), order;
        order.reserve(n);
        vector<int> stk; stk.reserve(n);
        vector<int> itIdx(n+1, 0);
        parent[c] = -1;
        stk.push_back(c);
        while(!stk.empty()){
            int u = stk.back();
            if(itIdx[u] < (int)g[u].size()){
                int v = g[u][itIdx[u]++];
                if(v==parent[u]) continue;
                parent[v]=u; stk.push_back(v);
            }else{
                order.push_back(u);
                stk.pop_back();
            }
        }
        static vector<ll> subsum, sumBest, valBest;
        static vector<int> height;
        if((int)subsum.size()!=n+1){
            subsum.assign(n+1,0);
            sumBest.assign(n+1,0);
            valBest.assign(n+1,0);
            height.assign(n+1,0);
        }
        for(int u: order){
            ll s = w[u];
            int h = 0;
            for(int v: g[u]){
                if(v==parent[u]) continue;
                s += subsum[v];
                h = max(h, height[v]+1);
            }
            subsum[u]=s; height[u]=h;
            if(h==0){
                sumBest[u]=s; valBest[u]=s;
            }else{
                ll bestShift = LLONG_MIN, ch_s=0, ch_v=0;
                for(int v: g[u]){
                    if(v==parent[u]) continue;
                    if(height[v]+1==h){
                        ll cand = valBest[v] + sumBest[v];
                        if(cand > bestShift){
                            bestShift=cand; ch_s=sumBest[v]; ch_v=valBest[v];
                        }
                    }
                }
                sumBest[u]=s + ch_s;
                valBest[u]=s + bestShift;
            }
        }

        // g_center = tổng trọng số toàn cây
        ll g_center = subsum[c];

        // Thu thập thông tin theo từng nhánh (hàng xóm của c)
        struct Item{ int id; int h; ll sumB, valB; };
        vector<Item> items;
        for(int x: g[c]){
            items.push_back({x, height[x], sumBest[x], valBest[x]});
        }
        // tìm hai độ sâu lớn nhất
        int h1=-1, h2=-1;
        for(auto &it: items){
            if(it.h > h1){ h2=h1; h1=it.h; }
            else if(it.h > h2){ h2=it.h; }
        }
        // các ứng viên cho mỗi phía
        vector<Item> A,B;
        for(auto &it: items){
            if(it.h==h1) A.push_back(it);
            if(it.h==h2) B.push_back(it);
        }

        auto pick_max = [](const vector<Item>& vec, bool useValPlusSum, int banId=-1)->Item{
            Item best{-1,-1,0,0};
            long long bestScore = LLONG_MIN;
            for(auto &it: vec){
                if(it.id==banId) continue;
                long long sc = useValPlusSum ? (it.valB + it.sumB) : it.valB;
                if(sc > bestScore){
                    bestScore = sc; best = it;
                }
            }
            return best;
        };

        ll ans = 0;

        // hướng 1: đi từ phía có h1 trước rồi qua tâm, rồi tới phía h2
        {
            Item L = pick_max(A, /*useValPlusSum=*/false);
            Item R = pick_max(B, /*useValPlusSum=*/true, (h1==h2?L.id:-1)); // nếu bằng nhau, tránh trùng nhánh
            if(L.id!=-1 && R.id!=-1){
                ll L1 = h1;
                ll S = L.valB + (L1+1)*g_center + (L1+1)*R.sumB + R.valB + R.sumB;
                ans = max(ans, S);
            }
        }
        // hướng 2: đảo lại
        {
            Item R = pick_max(B, /*useValPlusSum=*/false);
            Item L = pick_max(A, /*useValPlusSum=*/true, (h1==h2?R.id:-1));
            if(L.id!=-1 && R.id!=-1){
                ll L2 = h2;
                ll S = R.valB + (L2+1)*g_center + (L2+1)*L.sumB + L.valB + L.sumB;
                ans = max(ans, S);
            }
        }
        return ans;
    };

    long long answer = 0;
    if(dlen % 2 == 0){
        // tâm là đỉnh
        int c = path[dlen/2];
        answer = solve_center_vertex(c);
    }else{
        // tâm là cạnh (u, v)
        int u = path[dlen/2], v = path[dlen/2 + 1];
        CompRes A = solve_component(u, v);
        CompRes B = solve_component(v, u);
        // hai hướng
        long long S1 = A.valBest + (ll)A.height * B.sumBest + B.valBest;
        long long S2 = B.valBest + (ll)B.height * A.sumBest + A.valBest;
        answer = max(S1, S2);
    }

    cout << answer << "\n";
    return 0;
}