Language: C++

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

typedef pair<int, int> pii;
typedef long long ll;
typedef vector<int> vi;

#define pb push_back
#define eb emplace_back
#define mp make_pair
#define fi first
#define se second
#define rep(i,n) rep2(i,0,n)
#define rep2(i,m,n) for(int i=m;i<(n);i++)
#define ALL(c) (c).begin(),(c).end()
#define dump(x) cout << #x << " = " << (x) << endl

const ll INF = 1e18;

class unionfind {
	vector<int> par, rank;

public:
	void init(int n) {
		par.resize(n);
		rank.resize(n);

		for (int i = 0; i < n; i++) {
			par[i] = i;
			rank[i] = 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 (rank[x] < rank[y]) par[x] = y;
		else {
			par[y] = x;
			if (rank[x] == rank[y]) ++rank[x];
		}
	}

	bool same(int x, int y) { return (find(x) == find(y)); }
} uf;

const int MN = 200010;

using P = pair<ll, pii>;

vector<pii> g[MN], g2[MN];
int col[MN];
ll X[MN];
vector<P> dp[MN], res[MN];
P mt[MN];

vector<P> mergevec(vector<P> a, vector<P> b) {
	for (auto p : b) {
		a.eb(p);
	}
	sort(ALL(a));

	set<int> st;
	vector<P> res;

	for (auto p : a) {
		if (res.size() == 2) {
			break;
		}

		if (!st.count(p.se)) {
			st.insert(p.se);
			res.eb(p);
		}
	}

	return res;
}

void dfs_up(int v, int p) {
	dp[v].eb(X[v], col[v]);

	for (auto e : g[v]) {
		if (e.fi != p) {
			g2[v].eb(e);
			dfs_up(e.fi, v);

			vector<P> nx;

			for (auto t : dp[e.fi]) {
				nx.eb(t.fi + e.se, t.se);
			}

			dp[v] = mergevec(dp[v], nx);
		}
	}
	/*
	cerr << v << ", dpsize = " << dp[v].size() << endl;
	for (P pt : dp[v]) {
		cerr << pt.fi << " " << pt.se << endl;
	}*/
}

void dfs_down(int v, int p, vector<P> cand_parent) {
	vector<vector<P>> vl(g2[v].size() + 1), vr(g2[v].size() + 1);

	rep(i, g2[v].size()) {
		vl[i+1] = vl[i];
		vector<P> nx;
		int to = g2[v][i].fi;
		ll ct = g2[v][i].se;
		for (P t : dp[to]) {
			nx.eb(t.fi + ct, t.se);
		}
		vl[i+1] = mergevec(vl[i+1], nx);
	}

	for (int i = (int)g2[v].size() - 1; i >= 0; --i) {
		vr[i] = vr[i+1];
		vector<P> nx;
		int to = g2[v][i].fi;
		ll ct = g2[v][i].se;
		for (P t : dp[to]) {
			nx.eb(t.fi + ct, t.se);
		}
		vr[i] = mergevec(vr[i], nx);		
	}

	vector<P> cand = dp[v];
	cand = mergevec(cand, cand_parent);
	res[v] = cand;

	rep(i, g2[v].size()) {
		vector<P> cand_nx;

		int to = g2[v][i].fi;
		ll ct = g2[v][i].se;

		cand_nx = mergevec(cand_nx, vl[i]);
		cand_nx = mergevec(cand_nx, vr[i+1]);
		cand_nx = mergevec(cand_nx, cand_parent);
		cand_nx = mergevec(cand_nx, {P(X[v], col[v])});

		for (auto &t : cand_nx) {
			t.fi += ct;
		}

		dfs_down(to, v, cand_nx);
	}
}

int main() {
	int N; cin >> N;

	rep(i, N) cin >> X[i];

	rep(i, N-1) {
		int a, b, c;
		cin >> a >> b >> c;
		--a; --b;
		g[a].eb(b, c);
		g[b].eb(a, c);
	}

	uf.init(N);
	rep(i, N) {
		col[i] = i;
	}

	ll ans = 0;

	while (true) {
		bool one = true;
		for (int i = 1; i < N; ++i) {
			if (col[i] != col[0]) {
				one = false;
			}
		}
		if (one) {
			break;
		}

		rep(i, N) {
			dp[i].clear();
			g2[i].clear();
			mt[i] = mp(INF, -1);
		}

		dfs_up(0, -1);
		dfs_down(0, -1, {P(INF, -1)});

		rep(i, N) {
			bool found = false;
			for (P p : res[i]) {
				if (p.se != -1 && p.se != col[i]) {
					mt[col[i]] = min(mt[col[i]], mp(p.fi + X[i], p.se));
					found = true;
					break;
				}
			}
			assert(found);
		}

		rep(i, N) if (mt[i] != mp(INF, -1)) {
			//cerr << "mate " << i << " " << mt[i].se << endl;
			int u = i, v = mt[i].se;
			ll cost = mt[i].fi;
			if (!uf.same(u, v)) {
				uf.unite(u, v);
				ans += cost;
			}
		}

		rep(i, N) {
			col[i] = uf.find(col[i]);
		}
	}

	cout << ans << endl;

	return 0;
}

Compiler output:
a.cc: In function 'std::vector<std::pair<long long int, std::pair<int, int> > > mergevec(std::vector<std::pair<long long int, std::pair<int, int> > >, std::vector<std::pair<long long int, std::pair<int, int> > >)':
a.cc:78:30: error: no matching function for call to 'std::set<int>::count(std::pair<int, int>&)'
   78 |                 if (!st.count(p.se)) {
      |                      ~~~~~~~~^~~~~~
In file included from /usr/include/c++/14/set:63,
                 from /usr/include/x86_64-linux-gnu/c++/14/bits/stdc++.h:158,
                 from a.cc:1:
/usr/include/c++/14/bits/stl_set.h:756:9: note: candidate: 'template<class _Kt> decltype (((const std::set<_Key, _Compare, _Alloc>*)this)->std::set<_Key, _Compare, _Alloc>::_M_t._M_count_tr(__x)) std::set<_Key, _Compare, _Alloc>::count(const _Kt&) const [with _Key = int; _Compare = std::less<int>; _Alloc = std::allocator<int>]'
  756 |         count(const _Kt& __x) const
      |         ^~~~~
/usr/include/c++/14/bits/stl_set.h:756:9: note:   template argument deduction/substitution failed:
/usr/include/c++/14/bits/stl_set.h: In substitution of 'template<class _Kt> decltype (((const std::set<int>*)this)->std::set<int>::_M_t.std::_Rb_tree<int, int, std::_Identity<int>, std::less<int>, std::allocator<int> >::_M_count_tr<_Kt, _Req>(__x)) std::set<int>::count(const _Kt&) const [with _Kt = std::pair<int, int>]':
a.cc:78:16:   required from here
   78 |                 if (!st.count(p.se)) {
      |                      ~~~~~~~~^~~~~~
/usr/include/c++/14/bits/stl_set.h:757:37: error: no matching function for call to 'std::_Rb_tree<int, int, std::_Identity<int>, std::less<int>, std::allocator<int> >::_M_count_tr(const std::pair<int, int>&) const'
  757 |         -> decltype(_M_t._M_count_tr(__x))
      |                     ~~~~~~~~~~~~~~~~^~~~~
In file included from /usr/include/c++/14/map:62,
                 from /usr/include/x86_64-linux-gnu/c++/14/bits/stdc++.h:152:
/usr/include/c++/14/bits/stl_tree.h:1311:9: note: candidate: 'template<class _Kt, class _Req> std::_Rb_tree<_Key, _Val, _KeyOfValue, _Compare, _Alloc>::size_type std::_Rb_tree<_Key, _Val, _KeyOfValue, _Compare, _Alloc>::_M_count_tr(const _Kt&) const [with _Req = _Kt; _Key = int; _Val = int; _KeyOfValue = std::_Identity<int>; _Compare = std::less<int>; _Alloc = std::allocator<int>]'
 1311 |         _M_count_tr(const _Kt& __k) const
      |         ^~~~~~~~~~~
/usr/include/c++/14/bits/stl_tree.h:1311:9: note:   template argument deduction/substitution failed:
In file included from /usr/include/c++/14/string:49,
                 from /usr/include/c++/14/bitset:52,
                 from /usr/include/x86_64-linux-gnu/c++/14/bits/stdc++.h:52:
/usr/include/c++/14/bits/stl_function.h: In substitution of 'template<class _Func, class _SfinaeType> using std::__has_is_transparent_t = typename std::__has_is_transparent<_Func, _SfinaeType>::type [with _Func = std::less<int>; _SfinaeType = std::pair<int, int>]':
/usr/include/c++/14/bits/stl_tree.h:1309:9:   required by substitution of 'template<class _Kt> decltype (((const std::set<int>*)this)->std::set<int>::_M_t.std::_Rb_tree<int, int, std::_Identity<int>, std::less<int>, std::allocator<int> >::_M_count_tr<_Kt, _Req>(__x)) std::set<int>::count(const _Kt&) const [with _Kt = std::pair<int, int>]'
 1309 |                typename _Req = __has_is_transparent_t<_Compare, _Kt>>
      |                ^~~~~~~~
a.cc:78:16:   required from here
   78 |                 if (!st.count(p.se)) {
      |                      ~~~~~~~~^~~~~~
/usr/include/c++/14/bits/stl_function.h:1427:11: error: no type named 'type' in 'struct std::__has_is_transparent<std::less<int>, std::pair<int, int>, void>'
 1427 |     using __has_is_transparent_t
      |           ^~~~~~~~~~~~~~~~~~~~~~
/usr/include/c++/14/bits/stl_set.h:750:7: note: candidate: 'std::set<_Key, _Compare, _Alloc>::size_type std::set<_Key, _Compare, _Alloc>::count(const key_type&) const [with _Key = int; _Compare = std::less<int>; _Alloc = std::allocator<int>; size_type = long unsigned int; key_type = int]'
  750 |       count(const key_type& __x) const
      |       ^~~~~
/usr/include/c++/14/bits/stl_set.h:750:29: note:   no known conversion for argument 1 from 'std::pair<int, int>' to 'const std::set<int>::key_type&' {aka 'const int&'}
  750 |       count(const key_type& __x) const
      |             ~~~~~~~~~~~~~~~~^~~
a.cc:79:34: error: no matching function for call to 'std::set<int>::insert(std::pair<int, int>&)'
   79 |                         st.insert(p.se);
      |                         ~~~~~~~~~^~~~~~
/usr/include/c++/14/bits/stl_set.h:568:9: note: candidate: 'template<class _InputIterator> void std::set<_Key, _Compare, _Alloc>::insert(_InputIterator, _InputIterator) [with _Key = int; _Compare = std::less<int>; _Alloc = std::allocator<int>]'
  568 |         insert(_InputIterator __first, _InputIterator __last)
      |         ^~~~~~
/usr/include/c++/14/bits/stl_set.h:568:9: note:   candidate expects 2 arguments, 1 provided
/usr/include/c++/14/bits/stl_set.h:511:7: note: candidate: 'std::pair<typename std::_Rb_tree<_Key, _Key, std::_Identity<_Tp>, _Compare, typename __gnu_cxx::__alloc_traits<_Alloc>::rebind<_Key>::other>::const_iterator, bool> std::set<_Key, _Compare, _Alloc>::insert(const value_type&) [with _Key = int; _Compare = std::less<int>; _Alloc = std::allocator<int>; typename std::_Rb_tree<_Key, _Key, std::_Identity<_Tp>, _Compare, typename __gnu_cxx::__alloc_traits<_Alloc>::rebind<_Key>::other>::const_iterator = std::_Rb_tree<int, int, std::_Identity<int>, std::less<int>, std::allocator<int> >::const_iterator; typename __gnu_cxx::__alloc_traits<_Alloc>::rebind<_Key>::other = std::allocator<int>; typename __gnu_cxx::__alloc_traits<_Alloc>::rebind<_Key> = __gnu_cxx::__alloc_traits<std::allocator<int>, int>::rebind<int>; typename _Alloc::value_type = int; value_type = int]'
  511 |       insert(const value_type& __x)
      |       ^~~~~~
/usr/include/c++/14/bits/stl_set.h:511:32: note:   no known conversion for argument 1 from 'std::pair<int, int>' to 'const std::set<int>::value_type&' {aka 'const int&'}
  511 |       insert(const value_type& __x)
      |              ~~~~~~~~~~~~~~~~~~^~~
/usr/include/c++/14/bits/stl_set.h:520:7: note: candidate: 'std::pair<typename std::_Rb_tree<_Key, _Key, std::_Identity<_Tp>, _Compare, typename __gnu_cxx::__alloc_traits<_Alloc>::rebind<_Key>::other>::const_iterator, bool> std::set<_Key, _Compare, _Alloc>::insert(value_type&&) [with _Key = int; _Compare = std::less<int>; _Alloc = std::allocator<int>; typename std::_Rb_tree<_Key, _Key, std::_Identity<_Tp>, _Compare, typename __gnu_cxx::__alloc_traits<_Alloc>::rebind<_Key>::other>::const_iterator = std::_Rb_tree<int, int, std::_Identity<int>, std::less<int>, std::allocator<int> >::const_iterator; typename __gnu_cxx::__alloc_traits<_Alloc>::rebind<_Key>::other = std::allocator<int>; typename __gnu_cxx::__alloc_traits<_Alloc>::rebind<_Key> = __gnu_cxx::__alloc_traits<std::allocator<int>, int>::rebind<int>; typename _Alloc::value_type = int; value_type = int]'
  520 |       insert(value_type&& __x)
      |       ^~~~~~
/usr/include/c++/14/bits/stl_set.h:520:27: note:   no known conversion for argument 1 from 'std::pair<int, int>' to 'std::set<int>::value_type&&' {aka 'int&&'}
  520 |       insert(value_type&& __x)
      |              ~~~~~~~~~~~~~^~~
/usr/include/c++/14/bits/stl_set.h:548:7: note: candidate: 'std::set<_Key, _Compare, _Alloc>::iterator std::set<_Key, _Compare, _Alloc>::insert(const_iterator, const value_type&) [with _Key = int; _Compare = std::less<int>; _Alloc = std::allocator<int>; iterator = std::_Rb_tree<int, int, std::_Identity<int>, std::less<int>, std::allocator<int> >::const_iterator; const_iterator = std::_Rb_tree<int, int, std::_Identity<int>, std::less<int>, std::allocator<int> >::const_iterator; value_type = int]'
  548 |       insert(const_iterator __position, const value_type& __x)
      |       ^~~~~~
/usr/include/c++/14/bits/stl_set.h:548:7: note:   candidate expects 2 arguments, 1 provided
/usr/include/c++/14/bits/stl_set.h:553:7: note: candidate: 'std::set<_Key, _Compare, _Alloc>::iterator std::set<_Key, _Compare, _Alloc>::insert(const_iterator, value_type&&) [with _Key = int; _Compare = std::less<int>; _Alloc = std::allocator<int>; iterator = std::_Rb_tree<int, int, std::_Identity<int>, std::less<int>, std::allocator<int> >::const_iterator; const_iterator = std::_Rb_tree<int, int, std::_Identity<int>, std::less<int>, std::allocator<int> >::const_iterator; value_type = int]'
  553 |       insert(const_iterator __position, value_type&& __x)
      |       ^~~~~~
/usr/include/c++/14/bits/stl_set.h:553:7: note:   candidate expects 2 arguments, 1 provided
/usr/include/c++/14/bits/stl_set.h:580:7: note: candidate: 'void std::set<_Key, _Compare, _Alloc>::insert(std::initializer_list<_Tp>) [with _Key = int; _Compare = std::less<int>; _Alloc = std::allocator<int>]'
  580 |       insert(initializer_list<value_type> __l)
      |       ^~~~~~
/usr/include/c++/14/bits/stl_set.h:580:43: note:   no known conversion for argument 1 from 'std::pair<int, int>' to 'std::initializer_list<int>'
  580 |       insert(initializer_list<value_type> __l)
      |              ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^~~
/usr/include/c++/14/bits/stl_set.h:600:7: note: candidate: 'std::set<_Key, _Compare, _Alloc>::insert_return_type std::set<_Key, _Compare, _Alloc>::insert(node_type&&) [with _Key = int; _Compare = std::less<int>; _Alloc = std::allocator<int>; insert_return_type = std::_Rb_tree<int, int, std::_Identity<int>, std::less<int>, std::allocator<int> >::insert_return_type; node_type = std::_Rb_tree<int, int, std::_Identity<int>, std::less<int>, std::allocator<int> >::node_type]'
  600 |       insert(node_type&& __nh)
      |       ^~~~~~
/usr/include/c++/14/bits/stl_set.h:600:26: note:   no known conversion for argument 1 from 'std::pair<int, int>' to 'std::set<int>::node_type&&' {aka 'std::_Rb_tree<int, int, std::_Identity<int>, std::less<int>, std::allocator<int> >::node_type&&'}
  600 |       insert(node_type&& __nh)
      |              ~~~~~~~~~~~~^~~~
/usr/include/c++/14/bits/stl_set.h:605:7: note: candidate: 'std::set<_Key, _Compare, _Alloc>::iterator std