题解 | 最大最小路

#include <bits/stdc++.h>
using namespace std;
struct node{
    long long w,fu;
};
long long n,aa,bb;
int find(vector<struct node> &a,int x){
    if(a[x].fu<0)return x;
    return a[x].fu=find(a,a[x].fu);
}
void unio(vector<struct node> &a,int x,int y){
    int xx=find(a,x),yy=find(a,y);
    if(a[xx].fu<=a[yy].fu){
        a[xx].fu+=a[yy].fu;
        a[yy].fu=xx;
    }else{
        a[yy].fu+=a[xx].fu;
        a[xx].fu=yy;
    }
}
bool cond1(int w){
    if(w<bb)return true;
    return false;
}
bool cond2(int w){
    if(w>aa)return true;
    return false;
}
bool cond3(int w){
    if(w>aa&&w<bb)return true;
    return false;
}
long long ljs(vector<struct node> a,vector<pair<int,int>> &uv,bool (*cond)(int x)){
    long long ans=0;
    for(auto [u,v]:uv){
        if(cond(a[u].w)&&cond(a[v].w))unio(a,u,v);
    }
    for(int i=1;i<=n;i++){
        if(a[i].fu<0){
            long long nn=abs(a[i].fu);
            ans+=nn*(nn-1)/2;
        }
    }
    return ans;
}
int main() {
    cin>>n>>aa>>bb;
    vector<struct node> a(n+1);
    vector<pair<int,int>> uv(n-1);
    for(int i=1;i<=n;i++){
        cin>>a[i].w;
        a[i].fu=-1;
    }
    for(int i=0;i<n-1;i++){
        cin>>uv[i].first>>uv[i].second;
    }
    long long ans,ans1,ans2,ans3;
    ans=n*(n-1)/2;
    ans1=ljs(a,uv,cond1);
    ans2=ljs(a,uv,cond2);
    ans3=ljs(a,uv,cond3);
    cout<<ans-ans1-ans2+ans3;
}