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

题目链接  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;
}