【CodeForces】915 F. Imbalance Value of a Tree 并查集

【题目】F. Imbalance Value of a Tree

【题意】给定n个点的带点权树，求所有路径极差的和。n,ai<=10^6

【算法】并查集

【题解】先计算最大值的和，按点权从小到大排序，每个点x和相邻的已访问点的点集形成的路径的最大值都是a[x]，因为已访问过的点点权较小，然后用并查集并起来。复杂度O(n log n)。

#include<cstdio>
#include<algorithm>
using namespace std;
const int N=1000010;
struct edge{int v,from;}e[N*2];
int n,a
,b
,fa
,first
,tot,sz
;
long long ans;
int find(int x){return fa[x]==x?x:fa[x]=find(fa[x]);}
void insert(int u,int v){tot++;e[tot].v=v;e[tot].from=first[u];first[u]=tot;}
bool cmp(int x,int y){return a[x]<a[y];}
void solve(){
for(int i=1;i<=n;i++)fa[i]=0,b[i]=i;
sort(b+1,b+n+1,cmp);
for(int z=1;z<=n;z++){
int x=b[z];
fa[x]=x;sz[x]=1;
for(int i=first[x];i;i=e[i].from)if(fa[e[i].v]){
int y=find(e[i].v);
ans+=1ll*sz[x]*sz[y]*a[x];
sz[x]+=sz[y];
fa[y]=x;
}
}
}
int main(){
scanf("%d",&n);
for(int i=1;i<=n;i++)scanf("%d",&a[i]);
for(int i=1;i<n;i++){
int u,v;
scanf("%d%d",&u,&v);
insert(u,v);insert(v,u);
}
solve();
for(int i=1;i<=n;i++)a[i]=-a[i];
solve();
printf("%lld",ans);
return 0;
}