题解 | 最大最小路

#include<bits/stdc++.h>
using namespace std;
using ll=long long;
struct UF{
    UF* root;
    ll rank=0;
    ll cnt=1;
};


UF* Root(UF* a){
    if(a->root!=a){
        a->root=Root(a->root);
    }
    return a->root;
};

void merge(UF* rx,UF* ry){
    if(rx==ry)return;
    if(rx->rank<ry->rank){
        rx->root=ry;
        ry->cnt+=rx->cnt;
    }else{
        ry->root=rx;
        rx->cnt+=ry->cnt;
        if(rx->rank==ry->rank){
            rx->rank++;
        }
    }
}



int main(){
    ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
    ll n,a,b;
    cin>>n>>a>>b;
    ll sum=n*(n+1)/2;

    vector<ll>w(n+1);
    vector<UF>nodeA(n+1);
    vector<UF>nodeB(n+1);
    vector<UF>nodeAB(n+1);

    for(ll i=1;i<=n;i++){
        cin>>w[i];
        nodeA[i].root=&nodeA[i];
        nodeB[i].root=&nodeB[i];
        nodeAB[i].root=&nodeAB[i];
    }

    for(ll i=1;i<=n-1;i++){
        ll u,v;
        cin>>u>>v;
        if(w[u]>a&&w[v]>a){
            merge(Root(&nodeA[u]),Root(&nodeA[v]));
        }
        if(w[u]<b&&w[v]<b){
            merge(Root(&nodeB[u]),Root(&nodeB[v]));
        }
        if(w[u]<b&&w[u]>a&&w[v]<b&&w[v]>a){
            merge(Root(&nodeAB[u]),Root(&nodeAB[v]));
        }
    }

    ll sum_A=0,sum_B=0,sum_AB=0;
    unordered_set<UF*>sA;
    unordered_set<UF*>sB;
    unordered_set<UF*>sAB;
    for(ll i=1;i<=n;i++){
        if(sA.find(Root(&nodeA[i]))==sA.end()){
            ll t=Root(&nodeA[i])->cnt;
            sum_A+=t*(t+1)/2;
            sA.insert(Root(&nodeA[i]));
        }
        if(sB.find(Root(&nodeB[i]))==sB.end()){
            ll t=Root(&nodeB[i])->cnt;
            sum_B+=t*(t+1)/2;
            sB.insert(Root(&nodeB[i]));
        }
        if(sAB.find(Root(&nodeAB[i]))==sAB.end()){
            ll t=Root(&nodeAB[i])->cnt;
            sum_AB+=t*(t+1)/2;
            sAB.insert(Root(&nodeAB[i]));
        }
    }
    cout<<sum-sum_A-sum_B+sum_AB;
    
    return 0;
}