題目鏈接 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;
}

Codeforces 915F Imbalance Value of a Tree（並查集）