[可撤销并查集] Codeforces 1444C Team-Building

题目大意

给定一张 \(n\) 个点 \(m\) 条边的无向图，没有自环重边。
每一个结点都在一个颜色的组中，共有 \(k\) 组，可能存在某组为空。
求选出两组点，使它们能构成二分图的方案数。

题解

我们知道可以使用扩展域并查集来判二分图。即若存在边 \((u,v)\)，则把 \(u\) 和 \(v+n\) 所在的集合合并，把 \(u+n\) 和 \(v\) 所在的集合合并。若存在 \(u\) 和 \(u+n\) 在同一集合中，则构成奇环，不是二分图。

先把所有连接相同颜色的点的边加入并查集，分别判每种颜色的点构成的图是否是二分图。设有 \(a\) 种颜色的点自身无法构成二分图，那么还需考虑的颜色对数量为 \(\frac{1}{2}(n-a)\times(n-a-1)\)。

然后把所有连接不同颜色的点的边按两个点的颜色顺序排序，保证端点颜色相同的边相邻，一起处理。同时要忽略掉不能构成二分图的颜色。把每组端点颜色相同的边加入并查集，判这两个颜色的所有点能否构成二分图。若不能，答案减1。考虑完当前组边后，撤销当前组边合并的集合，然后考虑下一组边，所以需要使用可撤销并查集。时间复杂度 \(O(m\log m+m\log n)\)。

Code

#include <bits/stdc++.h>
using namespace std;

#define RG register int
#define LL long long

template<typename elemType>
inline void Read(elemType &T){
    elemType X=0,w=0; char ch=0;
    while(!isdigit(ch)) {w|=ch=='-';ch=getchar();}
    while(isdigit(ch)) X=(X<<3)+(X<<1)+(ch^48),ch=getchar();
    T=(w?-X:X);
}

template<size_t N>
struct UFS{
    int S[N],Rank[N];
    pair<int,int> stk[N*2];
    int top;

    void init(int n){
        top=0;
        for(int i=1;i<=n;++i)
            S[i]=i,Rank[i]=0;
    }
    int find(int u){
        while(u^S[u]) u=S[u];
        return u;
    }
    void merge_set(int u,int v){
        if((u=find(u))==(v=find(v))) return;
        if(Rank[u]<=Rank[v]){
            stk[++top]=make_pair(u,S[u]);
            S[u]=v;
            if(Rank[u]==Rank[v]){
                stk[++top]=make_pair(-v,Rank[v]);
                ++Rank[v];
            }
        }else{
            stk[++top]=make_pair(v,S[v]);
            S[v]=u;
        }
    }
    void undo(){
        int p=stk[top].first,u=stk[top].second;--top;
        if(p<0){
            Rank[-p]=u;
            p=stk[top].first,u=stk[top].second;--top;
            S[p]=u;
        }
        else S[p]=u;
    }
};

const int maxn=500010;
UFS<maxn*2> S;
vector<pair<int,int> > edge,data,banEdge;
int belong[maxn];
bool ban[maxn];
int N,M,K,banNum=0;

bool cmp(pair<int,int> A,pair<int,int> B){
    if(belong[A.first]==belong[B.first])
        return belong[A.second]<belong[B.second];
    return belong[A.first]<belong[B.first];
}

int main(){
    Read(N);Read(M);Read(K);
    S.init(N<<1);
    if(M==0){
        cout<<(LL)K*(K-1)/2<<endl;
        return 0;
    }
    S.init(N*2);
    for(int i=1;i<=N;++i)
        Read(belong[i]);
    int x,y;
    for(int i=1;i<=M;++i){
        int u,v;
        Read(u);Read(v);
        data.push_back(make_pair(u,v));
        if(belong[u]==belong[v]){
            S.merge_set(u,v+N);
            S.merge_set(u+N,v);
        }
    }
    int preTop=S.top;
    for(int u=1;u<=N;++u)
        if(S.find(u)==S.find(u+N))
            ban[belong[u]]=true;
    for(int i=1;i<=K;++i)
        if(ban[i]) ++banNum;
    for(auto e:data){
        int u=e.first,v=e.second;
        if(ban[belong[u]] || ban[belong[v]]) continue;
        if(belong[u]==belong[v]) continue;
        if(belong[u]>belong[v]) swap(u,v);
        edge.push_back(make_pair(u,v));
    }
    sort(edge.begin(),edge.end(),cmp);
    LL Ans=(LL)(K-banNum)*(LL)(K-banNum-1)>>1;
    int pre=0,pu=0,pv=0;
    for(int i=0;i<edge.size();++i){
        int u=edge[i].first,v=edge[i].second;
        if(belong[u]!=pu || belong[v]!=pv){
            while(S.top>preTop) S.undo();
            pre=i;pu=belong[u];pv=belong[v];
        }
        S.merge_set(u,v+N);
        S.merge_set(u+N,v);
        if(S.find(u)==S.find(u+N) || S.find(v)==S.find(v+N)){
            if(belong[u]>belong[v]) swap(u,v);
            banEdge.push_back(make_pair(belong[u],belong[v]));
        }
    }
    sort(banEdge.begin(),banEdge.end());
    banEdge.erase(unique(banEdge.begin(),banEdge.end()),banEdge.end());
    Ans-=banEdge.size();
    printf("%I64d\n",Ans);

    return 0;
}