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;
    }
}
using a3i = array<int, 3>;
void bucket_sort(vector<a3i> &v, int dir)
{
    vector<a3i> ori(v);
    vector<int> lo[1024], hi[1024];
    int mask = (1 << 10) - 1;
    for (int i = 0; i < v.size(); ++i)
        lo[v[i][0] & mask].push_back(i);
    for (int i = 0; i < 1024; ++i)
        for (auto idx : lo[i])
            hi[v[idx][0] >> 10].push_back(idx);

    int pos = 0;
    for (int i = 0; i < 1024; ++i)
        for (auto idx : hi[i])
            v[pos++] = ori[idx];
    if (dir == 1)
        reverse(v.begin(), v.end());
}

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);
    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());
    bucket_sort(es, 0);
    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());
    bucket_sort(es, 1);
    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;
};