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;
vi g[1000003];
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;

    for (int i = 1; i < n; ++i) {
        int a, b;
        cin >> a >> b;
        g[a].push_back(b), g[b].push_back(a);
    }
    init_dsu();
    sort(a.begin(), a.end());
    set<int> exist;
    for (int i = 0; i < n; ++i) {
        int u = a[i].second;
        ll w = a[i].first;
        exist.insert(u);
        for (auto v : g[u])
            if (exist.count(v)) {
                int gp = find(v);
                int gsize = cnt[gp];
                int usize = cnt[u];
                join(gp, u);
                maxv += w * gsize * usize;
            }
    }

    init_dsu();
    exist.clear();
    sort(a.rbegin(), a.rend());
    for (int i = 0; i < n; ++i) {
        int u = a[i].second;
        ll w = a[i].first;
        exist.insert(u);
        for (auto v : g[u]) {
            if (exist.count(v)) {
                int gp = find(v);
                int gsize = cnt[gp];
                int usize = cnt[u];
                join(gp, u);
                minv += w * gsize * usize;
            }
        }
    }

    cout << maxv - minv << endl;

    return 0;
};