Codeforces 1444C 可撤销并查集

题意

传送门 Codeforces 1444C Team-Building

题解

图中节点划分为不相交的点集 $S$，求满足所构成子图为二分图的点集对 $(S_1,S_2)$ 数量。

若点集 $S$ 构成的子图存在奇环，则它与任意其他点集构成的子图都不可能是二分图。讨论满足二分图判定的点集即可。

若两点间存在一条边，则两点划分至不同的集合，可以通过并查集维护这样的关系。依次考虑点集 $S$，处理与它存在连边的点集 $T$，判断 $S,T$ 构成的子图是否为二分图即可。可以用可撤销并查集维护这样的关系。

总时间复杂度 $O(m\log n)$。

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
constexpr int MAXN = 5E5 + 5;
int N, M, K, C[MAXN];
struct edge
{
    int u, v, id;
};
vector<edge> es[MAXN], es2[MAXN];

struct DSU
{
    int par[MAXN * 2], rnk[MAXN * 2], n;
    int stk[MAXN * 2][2], top;
    void init(int _n)
    {
        n = _n, top = 0;
        for (int i = 0; i < n; ++i)
            par[i] = i, rnk[i] = 0;
    }
    int find(int x) { return par[x] == x ? x : find(par[x]); }
    bool same(int x, int y) { return find(x) == find(y); }
    void merge(int x, int y)
    {
        x = find(x), y = find(y);
        if (x == y)
            return;
        if (rnk[x] > rnk[y])
            swap(x, y);
        int d = rnk[x] == rnk[y];
        stk[top][0] = x, stk[top++][1] = d;
        par[x] = y, rnk[y] += d;
    }
    void undo(int t)
    {
        while (top > t)
        {
            --top;
            int x = stk[top][0], d = stk[top][1];
            rnk[par[x]] -= d, par[x] = x;
        }
    }
} dsu;

int main()
{
    ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
    cin >> N >> M >> K;
    for (int i = 0; i < N; ++i)
        cin >> C[i], --C[i];
    for (int i = 0; i < M; ++i)
    {
        int u, v;
        cin >> u >> v;
        --u, --v;
        if (C[u] == C[v])
            es[C[u]].push_back({u, v, C[v]});
        else
        {
            if (C[u] < C[v])
                swap(u, v);
            es2[C[u]].push_back({u, v, C[v]});
        }
    }
    
    ll res = 0;
    dsu.init(2 * N);
    set<int> st;
    for (int i = 0; i < K; ++i)
    {
        bool ok = 1;
        for (auto &e : es[i])
        {
            int u = e.u, v = e.v;
            dsu.merge(u, N + v), dsu.merge(N + u, v);
            if (dsu.same(u, v) || dsu.same(N + u, N + v))
                ok = 0;
        }
        if (ok == 0)
            continue;
        set<int> st2;
        sort(es2[i].begin(), es2[i].end(), [](const edge &x, const edge &y)
             { return x.id < y.id; });
        for (int l = 0, r = 0, n = es2[i].size();l < n;)
        {
            while (r < n && es2[i][l].id == es2[i][r].id)
                ++r;
            int bg = dsu.top, g = es2[i][l].id;
            bool ok = 1;
            for (int j = l; j < r; ++j)
            {
                int u = es2[i][j].u, v = es2[i][j].v;
                dsu.merge(u, N + v), dsu.merge(N + u, v);
                if (dsu.same(u, v) || dsu.same(N + u, N + v))
                    ok = 0;
            }
            if (ok == 0)
                if (st.count(g))
                    st.erase(g), st2.insert(g);
            dsu.undo(bg);
            l = r;
        }
        res += st.size();
        for (int g : st2)
            st.insert(g);
        st.insert(i);
    }
    cout << res << '\n';
    return 0;
}