CF1444C Team-Building 可撤销种类并查集

• 思路：

1. 该题很容易想到，暴力枚举子图，然后二分图判定，统计答案即可，但是复杂度太高。
2. 正难则反，我们考虑用所有方案减去不合法方案。
3. 不和法方案分为两种：
   - 如果一个环为奇环，并且只有一种颜色，那么所有和他构成的偏序集都是不合法的
   - 如果一个环为奇环，并且环内由两种颜色构成，那么由这两种颜色构成的偏序集是不合法的
4. 对于第一种情况，利用并查集可以很快的判断出来。
5. 第二种情况只需要先将图进行一些处理，先将所有相同颜色的边利用并查集构成连通块，这个时候利用的是连接相同颜色的边。然后再将连接不同颜色的边分组。分组之后，每次按组利用并查集找出奇环，这个时候找出的奇环一定是两个颜色的奇环。每次统计完一组的答案之后，需要消除这组边的影响，所以利用可撤销并查集撤销之前的操作即可。
   代码如下：

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

#define ll long long
#define mem(a, b) memset(a,b,sizeof(a))
#define sz(a) (int)a.size()
#define INF 0x3f3f3f3f
#define DNF 0x7f
#define DBG printf("this is a input\n")
#define fi first
#define se second
#define mk(a, b) make_pair(a,b)
#define pb push_back
#define LF putchar('\n')
#define SP putchar(' ')
#define p_queue priority_queue
#define CLOSE ios::sync_with_stdio(0); cin.tie(0)
#define sz(a) (int)a.size()
#define pii pair <int,int>
template<typename T>
void read(T &x) {x = 0;char ch = getchar();ll f = 1;while(!isdigit(ch)){if(ch == '-')f *= -1;ch = getchar();}while(isdigit(ch)){x = x * 10 + ch - 48; ch = getchar();}x *= f;}
template<typename T, typename... Args>
void read(T &first, Args& ... args) {read(first);read(args...);}
template<typename T>
void write(T arg) {T x = arg;if(x < 0) {putchar('-'); x =- x;}if(x > 9) {write(x / 10);}putchar(x % 10 + '0');}
template<typename T, typename ... Ts>
void write(T arg, Ts ... args) {write(arg);if(sizeof...(args) != 0) {putchar(' ');write(args ...);}}
using namespace std;

ll gcd(ll a, ll b) {
    return b == 0 ? a : gcd(b, a % b);
}

ll lcm(ll a, ll b) {
    return a / gcd(a, b) * b;
}

const int N = 1000005;
map <pii , int> ma;
vector <int> ve[N];
int n , m , k;
int c[N], a[N], b[N];
int f[N];
struct Undo_Dsu
{
    stack <pii> st ;
    int fa[N] , siz[N] ;
    void init()
    {
        while(!st.empty())  st.pop() ;
        for (int i = 1 ; i <= n * 2 + 1; i ++)  fa[i] = i , siz[i] = 1 ;
    }
    int findroot(int x)
    {
        return x == fa[x] ? x : findroot(fa[x]) ;
    }
    bool merge(int u , int v)
    {
        int fax = findroot(u) , fay = findroot(v);
        if(fax == fay)  return false ;
        if (siz[fax] > siz[fay])  swap(fax, fay) , swap(u , v) ;
        fa[fax] = fay , siz[fay] += siz[fax] ;
        st.push({fax , fay}) ;
        return true ;
    }
    void undo() //撤销一次最新的合并
    {
        pii no = st.top();
        fa[no.fi] = no.fi ;
        siz[no.se] -= siz[no.fi] ;
        st.pop() ;
    }
} dsu ;
int main ()
{
    ll ans = 0;
    read (n, m, k);
    dsu.init();
    for (int i = 1 ; i <= n ; i ++)
        read (c[i]);
    for (int i = 1 ; i <= m ; i ++)
        read (a[i] , b[i]);
    for (int i = 1 ; i <= m ; i ++)
    {
        if (c[a[i]] == c[b[i]])
        {
            int fau = dsu.findroot (a[i]);
            int fav = dsu.findroot (b[i]);
            if (fau == fav) f[c[a[i]]] = 1;
            else dsu.merge(a[i], b[i] + n) , dsu.merge(b[i], a[i] + n); //用于奇环的判定，可以参考“LuoGu 关押罪犯”
        }
    }
    //按连接组将边分类，连接奇环的不予考虑
    int type = 0;
    for (int i = 1 ; i <= m ; i ++)
    {
        if (c[a[i]] != c[b[i]] && !f[c[a[i]]] && !f[c[b[i]]])
        {
            int lc = c[a[i]] , rc = c[b[i]];
            if (lc > rc) swap (lc , rc);
            int id = ma[mk(lc, rc)];
            if (!id)
                id = ma[mk(lc, rc)] = ++ type;
            ve[id].pb (i);
        }
    }
    //统计答案
    for (auto x : ma)
    {
        int w = x.se;
        int now = sz(dsu.st);
        for (auto vv : ve[w])
        {
            int uid = a[vv], vid = b[vv];
            int fau = dsu.findroot(uid), fav = dsu.findroot (vid);
            if (fau == fav)
            {
                -- ans;
                break;
            }
            dsu.merge (uid, vid + n);
            dsu.merge (vid, uid + n);
        }
        while (sz(dsu.st) > now) dsu.undo();
    }
    int cnt = 0;
    for (int i = 1 ; i <= k ; i ++) if (!f[i]) cnt ++;
    ans += 1ll*cnt*(cnt-1)/2;
    write (ans), LF;
}