题目链接:Codeforces - Imbalance Value of a Tree
由原来的式子,可以转化为:ΣΣ(max()-min())=ΣΣmax()-ΣΣmin()
所以我们分别求max和min即可。
怎么求呢?考虑并查集合并途中计算答案。如果计算最大值。
我们按照点从小到大合并。这样统计点的size就可以了。
计算最小值同理。
AC代码:
#pragma GCC optimize("-Ofast","-funroll-all-loops")
#include<bits/stdc++.h>
#define int long long
using namespace std;
const int N=1e6+10;
int n,f[N],sz[N],res,a[N];
struct node{int u,v,w1,w2;}t[N];
int cmp1(node a,node b){return a.w2<b.w2;}
int cmp2(node a,node b){return a.w1>b.w1;}
int find(int x){return x==f[x]?x:f[x]=find(f[x]);}
signed main(){
ios::sync_with_stdio(false),cin.tie(nullptr);
cin>>n;
for(int i=1;i<=n;i++) cin>>a[i];
for(int i=1,x,y;i<n;i++) cin>>x>>y,t[i]={x,y,min(a[x],a[y]),max(a[x],a[y])};
sort(t+1,t+n,cmp1);
for(int i=1;i<=n;i++) f[i]=i,sz[i]=1;
for(int i=1;i<n;i++){
int x=find(t[i].u),y=find(t[i].v);
if(x==y) continue;
f[x]=y; res+=t[i].w2*sz[x]*sz[y]; sz[y]+=sz[x];
}
for(int i=1;i<=n;i++) f[i]=i,sz[i]=1;
sort(t+1,t+n,cmp2);
for(int i=1;i<n;i++){
int x=find(t[i].u),y=find(t[i].v);
if(x==y) continue;
f[x]=y; res-=t[i].w1*sz[x]*sz[y]; sz[y]+=sz[x];
}
cout<<res;
return 0;
}