Imbalance Value of a Tree CodeForces - 915F

链接

大意: 给定树, 求树上所有链上最大值最小值之差

817D的树上版本, 用并查集维护即可.

817D由于是链的情况并查集不必压缩路径即可达到均摊$O(n)$, 该题必须压缩, 复杂度$O(nlogn)$

#include <iostream>
#include <algorithm>
#include <cstdio>
#include <vector>
#define REP(i,a,n) for(int i=a;i<=n;++i)
#define pb push_back
#define x first
#define y second
using namespace std;
typedef long long ll;
typedef pair<int,int> pii;

const int N = 1e6+10, INF = 0x3f3f3f3f;
int n;
pii a[N];
vector<int> g[N];
ll ans;
int fa[N], sz[N];
int Find(int x) {return fa[x]==x?x:fa[x]=Find(fa[x]);}

void solve() {
	sort(a+1,a+1+n);
	REP(i,1,n) fa[i]=0,sz[i]=1;
	REP(i,1,n) {
		int x = a[i].y, w = a[i].x;
		fa[x] = x;
		for (int y:g[x]) if (fa[y]) {
			y = Find(y);
			ans += (ll)sz[x]*sz[y]*w;
			sz[x] += sz[y];
			fa[y] = x;
		}
	}
}

int main() {
	scanf("%d", &n);
	REP(i,1,n) scanf("%d", &a[i].x), a[i].y=i;
	REP(i,2,n) {
		int u, v;
		scanf("%d%d", &u, &v);
		g[u].pb(v),g[v].pb(u);
	}
	solve();
	REP(i,1,n) a[i].x=-a[i].x;
	solve();
	printf("%lld\n", ans);
}