题目链接 Imbalance Value of a Tree
题意 给定一棵树。求树上所有简单路径中的最大权值与最小权值的差值的和。
首先考虑求所有简单路径中的最大权值和。
对所有点按照权值大小升序排序,即若$a[i] < a[j]$,那么$i$排在$j$前面。
接下来开始依次处理。对于每个点$i$,寻找周围跟他连通并且权值比他小的点进行合并,并且累加答案。
整个过程用并查集维护。
#include <bits/stdc++.h>
using namespace std;
#define rep(i, a, b) for (int i(a); i <= (b); ++i)
#define dec(i, a, b) for (int i(a); i >= (b); --i)
#define MP make_pair
#define fi first
#define se second
typedef long long LL;
const int N = 1e6 + 10;
int a[N];
int sz[N];
int n;
int father[N];
int c[N];
LL ans = 0;
vector <int> v[N];
int getfather(int x){
return father[x] == x ? x : father[x] = getfather(father[x]);
}
bool cmp(const int &x, const int &y){
return a[x] < a[y];
}
void work(int x, int y, int z){
int fx = getfather(x);
int fy = getfather(y);
if (!fx || !fy) return;
ans += 1ll * z * sz[fx] * sz[fy];
father[fx] = fy;
sz[fy] += sz[fx];
sz[fx] = 0;
}
void solve(){
rep(i, 1, n) c[i] = i;
sort(c + 1, c + n + 1, cmp);
rep(i, 1, n) father[i] = 0, sz[i] = 0;
rep(i, 1, n){
int x = c[i];
father[x] = x;
sz[x] = 1;
for (auto u : v[x]){
work(x, u, a[x]);
}
}
}
int main(){
scanf("%d", &n);
rep(i, 1, n) scanf("%d", a + i);
rep(i, 2, n){
int x, y;
scanf("%d%d", &x, &y);
v[x].push_back(y);
v[y].push_back(x);
}
solve();
rep(i, 1, n) a[i] *= -1;
solve();
printf("%lld\n", ans);
return 0;
}