Language: C++

#include <iostream>
#include <vector>
#include <string>
#include <cmath>
#include <algorithm>
#include <utility>
#include <queue>
#include <set>
#include <map>
#include <deque>
#include <iomanip>
#include <cstdio>
#include <stack>

using namespace std;
typedef  long long ll;
typedef pair<int,int> PII;
typedef vector<int> VI;
typedef vector<VI> VVI;
#define  MP make_pair
#define  PB push_back
#define inf  1000000007
#define rep(i,n) for(int i=0;i<(int)(n);++i)

class UnionFind {
private:
    int sz;
    vector<int> par, nrank;
public:
    UnionFind(){}
    UnionFind(int node_size) : sz(node_size), par(sz), nrank(sz, 0){
        iota(par.begin(), par.end(), 0);
    }
    int find(int x){
        if(par[x] == x) return x;
        else return par[x] = find(par[x]);
    }
    void unite(int x,int y){
        x = find(x), y = find(y);
        if(x == y) return;
    	if(nrank[x] < nrank[y]) swap(x,y);
        par[y] = x;
        if(nrank[x] == nrank[y]) nrank[x]++;
    }
    bool same(int x,int y){
        return find(x) == find(y);
    }
};
 


pair<pair<ll,int>,pair<ll,int> > dp1[200010];
pair<pair<ll,int>,pair<ll,int> > dp2[200010];
vector<ll> x;
vector<int> q;
vector<vector<pair<int,ll> > >g;

pair<pair<ll,int>,pair<ll,int> > merge(pair<pair<ll,int>,pair<ll,int> >a,pair<pair<ll,int>,pair<ll,int> >b){
    vector<pair<ll,int> > s;
    s.PB(a.first);
    s.PB(a.second);
    s.PB(b.first);
    s.PB(b.second);
    sort(s.begin(),s.end());
    if(s[0].second==s[1].second){
        return MP(s[0],s[2]);
    }else{
        return MP(s[0],s[1]);
    }
}
pair<pair<ll,int>,pair<ll,int> > merge(pair<pair<ll,int>,pair<ll,int> >a,pair<ll,int> b){
    if(a.first.second==b.second){
        if(b.first>a.first.first){
            return a;
        }else{
            return MP(b,a.second);
        }
    }else{
        if(b.first<a.first.first){
            return MP(b,a.first);
        }else{
            if(b.first<a.second.first){
                return MP(a.first,b);
            }else{
                return a;
            }
        }
    }
}
pair<pair<ll,int>,pair<ll,int> > dfs(int a,int pre){
    pair<pair<ll,int>,pair<ll,int> > tmp;
    tmp.first = MP((1LL<<60),-1);
    tmp.second =MP((1LL<<60),-1);
    for(auto p:g[a]){
        if(p.first!=pre){
            pair<pair<ll,int>,pair<ll,int> > s = dfs(p.first,a);
            s.first.first += p.second;
            s.second.first += p.second;
            tmp = merge(tmp,s);
        }
    }
    dp1[a] = tmp;
    return merge(dp1[a],MP(x[a],q[a]));
}
void dfs2(int a,int pre,ll val){
    pair<pair<ll,int>,pair<ll,int> > tmp = dp1[a];
    if(pre!=-1){
        tmp = merge(tmp,MP(x[pre]+val,pre));
        auto s = dp2[pre];
        s.first.first += val;
        s.second.first += val;
        tmp = merge(tmp,s);
    }
    dp2[a] = tmp;
    for(auto p:g[a]){
        if(p.first!=pre){
            dfs2(p.first,a,p.second);
        }
    }
}

int main(){
    int n;
    cin >> n;
    x.resize(n);
    g.resize(n);
    rep(i,n)cin >> x[i];
    rep(i,n-1){
        int a,b;
        ll c;
        cin >> a >> b >> c;
        a--;b--;
        g[a].PB(MP(b,c));
        g[b].PB(MP(a,c));
    }
    UnionFind uf(n);
    q.resize(n);
    rep(i,n){
        q[i] = i;
    }
    ll ans = 0;
    while(1){
        dfs(0,-1);
        dfs2(0,-1,0);
        vector<pair<ll,int> >r(n,MP(1LL<<60,-1));
        rep(i,n){
            if(!uf.same(q[i],dp2[i].first.second)){
                r[q[i]] = min(r[q[i]],MP(dp2[i].first.first+x[i],dp2[i].first.second));
            }else{
                r[q[i]] = min(r[q[i]],MP(dp2[i].second.first+x[i],dp2[i].second.second));
            }
        }
        rep(i,n){
            if(r[i].second!=-1&&!uf.same(i,r[i].second)){
                //cerr << i << " " << r[i].second << " " << r[i].first << endl;
                ans += r[i].first;
                uf.unite(i,r[i].second);
            }
        }
        bool flag = 1;
        rep(i,n){
            q[i] = uf.find(i);
            if(q[0]!=q[i]){
                flag = 0;
            }
        }
        if(flag)break;
    }
    cout << ans << endl;
    return 0;
}

Compiler output:
a.cc: In constructor 'UnionFind::UnionFind(int)':
a.cc:32:9: error: 'iota' was not declared in this scope
   32 |         iota(par.begin(), par.end(), 0);
      |         ^~~~