【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[N],b[N],fa[N],first[N],tot,sz[N];
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;
}