CF915F TLE

// Problem: F. Imbalance Value of a Tree
// Contest: Codeforces - Educational Codeforces Round 36 (Rated for Div. 2)
// URL: https://codeforces.com/problemset/problem/915/F
// Memory Limit: 256 MB
// Time Limit: 4000 ms
//
// Powered by CP Editor (https://cpeditor.org)

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

template <class T> using mpq = priority_queue<T, vector<T>, greater<T>>;

using ll = long long;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
using vl = vector<ll>;
using vi = vector<int>;

int n;
vector<pii> a;

ll maxv, minv;

int fa[1000003], cnt[1000003];
void init_dsu()
{
    memset(fa, -1, sizeof(fa));
    fill_n(cnt, sizeof(cnt) / sizeof(int), 1);
}
int find(int x)
{
    if (fa[x] == -1)
        return x;
    return fa[x] = find(fa[x]);
}
void join(int x, int y)
{
    int px = find(x), py = find(y);
    if (px != py) {
        fa[px] = py;
        cnt[py] += cnt[px];
        cnt[px] = 0;
    }
}

int main(int argc, char **argv)
{
    cin >> n;
    a = vector<pii>(n);
    for (int i = 0; i < n; ++i)
        cin >> a[i].first, a[i].second = i + 1;

    vector<pii> edges(n - 1);
    using a3i = array<int, 3>;
    vector<a3i> es(n - 1);
    for (int i = 0; i < n - 1; ++i)
        cin >> edges[i].first >> edges[i].second;

    init_dsu();
    // sort(a.begin(), a.end());
    for (int i = 0; i < n - 1; ++i)
        es[i][1] = edges[i].first, es[i][2] = edges[i].second,
        es[i][0] = max(a[es[i][1] - 1].first, a[es[i][2] - 1].first);
    sort(es.begin(), es.end());
    for (int i = 0; i < n - 1; ++i) {
        int u = es[i][1], v = es[i][2];
        ll w = es[i][0];
        int up = find(u), vp = find(v);
        int usize = cnt[up], vsize = cnt[vp];
        join(up, vp);
        maxv += w * usize * vsize;
    }

    init_dsu();
    for (int i = 0; i < n - 1; ++i)
        es[i][1] = edges[i].first, es[i][2] = edges[i].second,
        es[i][0] = min(a[es[i][1] - 1].first, a[es[i][2] - 1].first);
    sort(es.rbegin(), es.rend());
    for (int i = 0; i < n - 1; ++i) {
        int u = es[i][1], v = es[i][2];
        ll w = es[i][0];

        int up = find(u), vp = find(v);
        int usize = cnt[up], vsize = cnt[vp];
        join(up, vp);
        minv += w * usize * vsize;
    }

    cout << maxv - minv << endl;

    return 0;
};