題目大意

給定一張 \(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;
}