CF915F Imbalance Value of a Tree

标程：

 1 #include<bits/stdc++.h>
 2 using namespace std;
 3 typedef long long ll;
 4 const int N=1000005;
 5 struct node{int u,v,w;}e[N];
 6 int a[N],f[N],sz[N],n;
 7 ll ans;
 8 int find(int x){return x==f[x]?x:f[x]=find(f[x]);}
 9 bool operator < (const node &A,const node &B) {return A.w<B.w;}
10 void solve()
11 {
12   for (int i=1;i<=n;i++) f[i]=i,sz[i]=1;
13   for (int i=1;i<n;i++) e[i].w=max(a[e[i].u],a[e[i].v]);
14   sort(e+1,e+n);
15   for (int i=1;i<n;i++)
16   {
17     int x=find(e[i].u),y=find(e[i].v);
18     if (x==y) continue;
19     ans+=(ll)e[i].w*sz[x]*sz[y];
20     f[x]=y;sz[y]+=sz[x];
21   }
22 }
23 int main()
24 {
25    scanf("%d",&n);
26    for (int i=1;i<=n;i++) scanf("%d",&a[i]);
27    for (int i=1;i<n;i++) scanf("%d%d",&e[i].u,&e[i].v);
28    solve();
29    for (int i=1;i<=n;i++) a[i]=-a[i];
30    solve();
31    printf("%lld\n",ans);
32    return 0;
33 }