求助 WA 29,没有路径压缩
查看原帖
求助 WA 29,没有路径压缩
237530
rzh123楼主2024/1/22 23:30
#include <bits/stdc++.h>
using namespace std;
constexpr int N=2e6+5;
int n,m,kk,bel[N];
array<basic_string<int>,N> g;
vector<int> grp[N];
bool bipa[N];
long long ans{0ll};
struct DSU{
    int fa[N],sz[N];
    stack<tuple<int,int,int> > stk;
    inline void init(int n){
        iota(fa+1,fa+n+1,1);
        fill_n(sz+1,n,1);
    }
    int find(int u)const{return fa[u]==u?u:find(fa[u]);}
    inline void merge(int u,int v){
        // printf("mrg %d %d\n",u,v);
        u=find(u),v=find(v);
        if(u==v) return;
        if(sz[u]>sz[v]) swap(u,v);
        stk.emplace(u,v,sz[v]);
        sz[v]+=sz[u],fa[u]=v;
    }
    inline void backup(){
        // printf("bkup\n");
        while(!stk.empty()) stk.pop();
    }
    inline void undo(){
        // printf("undo\n");
        while(!stk.empty()){
            auto t=stk.top(); stk.pop();
            fa[get<0>(t)]=get<0>(t);
            assert(sz[get<1>(t)]-sz[get<0>(t)]==get<2>(t));
            sz[get<1>(t)]=get<2>(t);
        }
    }
}dsu;
int main(){
    scanf("%d%d%d",&n,&m,&kk);
    vector<pair<int,int>> ee;
    for(int i{1};i<=n;++i) scanf("%d",&bel[i]),grp[bel[i]].emplace_back(i);
    for(int i{1};i<=m;++i){
        int u,v; cin>>u>>v;
        g[u]+=v,g[v]+=u;
        if(bel[u]>bel[v]) swap(u,v);
        if(bel[u]!=bel[v]) ee.emplace_back(u,v);
    }
    dsu.init(n<<1);
    for(int i{1};i<=kk;++i){
        bipa[i]=true;
        for(int u:grp[i]){
            for(int v:g[u]){
                if(bel[v]==i)
                    dsu.merge(u,v+n),dsu.merge(v,u+n);
            }
        }
        for(int u:grp[i]){
            if(dsu.find(u)==dsu.find(u+n)){bipa[i]=false;break;}
        }
    }
    int cc{0};
    for(int i{kk};i>=1;--i)
        if(bipa[i])
            ans+=cc,++cc;
    sort(ee.begin(),ee.end(),
        [](const pair<int,int> &ea,const pair<int,int> &eb){
            return make_pair(bel[ea.first],bel[ea.second])<make_pair(bel[eb.first],bel[eb.second]);
        }
    );  
    dsu.backup();
    assert(dsu.stk.empty());
    pair<int,int> lst{0,0};
    bool crtok{true};
    ee.emplace_back(0,0);
    for(const auto &p:ee){
        auto [u,v]=p;
        if(make_pair(bel[u],bel[v])!=lst){
            if(lst.first!=0&&bipa[bel[lst.first]]&&bipa[bel[lst.second]]&&!crtok) --ans;
            dsu.undo(),crtok=true,lst=make_pair(bel[u],bel[v]);
        }
        if(!u||!bipa[bel[u]]||!bipa[bel[v]]) continue;
        dsu.merge(u,v+n),dsu.merge(u+n,v);
        if(dsu.find(u)==dsu.find(u+n)||dsu.find(v)==dsu.find(v+n)) crtok=false;
    }
    printf("%lld\n",ans);
    return 0;
}
2024/1/22 23:30
加载中...