Submission #1603411

---

Source Code Expand

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

#define all(x)  x.begin(), x.end()
#define mp      make_pair
#define pii     pair<int, int>
#define pll     pair<long long, long long>
#define ll      long long

struct UnionFind {
        int n;
        vector<int> parent;
        vector<int> rank;
        vector<int> num;
        int find(int x) {
                if (parent[x] == x) return  x;
                return parent[x] = find(parent[x]);
        }
        UnionFind(int n_) {
                n = n_;
                parent.resize(n);
                for (int i = 0; i < n; i ++) parent[i] = i;
                rank.assign(n, 0);
                num.assign(n, 1);
        }
        void unite(int x, int y) {
                if ((x = find(x)) != (y = find(y))) {
                        if (rank[x] < rank[y]) {
                                parent[x] = y;
                                num[y] += num[x];
                        } else {
                                parent[y] = x;
                                if (rank[x] == rank[y]) rank[x] ++;
                                num[x] += num[y];
                        }
                        n --;
                }
        }
        bool same(int x, int y) { return find(x) == find(y); }
        int get() { return n; }
        int get(int x) { return num[find(x)]; }
};

vector<int> g[100000], rg[100000];

int color[100000];
int A, B, C, edge;

int dfs(int v, int c) {
        if (color[v] == -1) { 
                color[v] = c;
                if (c == 0) A ++;
                else if (c == 1) B ++;
                else C ++;
        } else if (color[v] != c) return -1;
        else return 0;
        for (int u : g[v]) { 
                edge ++;
                if (dfs(u, (c + 1) % 3) == -1) return -1;
        }
        for (int u : rg[v]) { 
                edge ++;
                if (dfs(u, (c - 1 + 3) % 3) == -1) return -1;
        }
        return 0;
}

int main() {
        int n, m;
        cin >> n >> m;
        UnionFind uf(n + 1);
        for (int i = 0; i < m; i ++) {
                int a, b;
                cin >> a >> b;
                a --, b --;
                g[a].push_back(b);
                rg[b].push_back(a);
                uf.unite(a, b);
        }
        memset(color, -1, sizeof color);
        long long ans = 0;
        for (int i = 0; i < n; i ++) if (!uf.same(i, n)) {
                long long num = uf.get(i);
                uf.unite(i, n); //used
                A = 0, B = 0, C = 0, edge = 0;
                if (dfs(i, 0) == -1) ans += num * num;
                else if (C == 0 || B == 0) ans += edge / 2;
                else ans += 1LL * A * B + 1LL * B * C + 1LL * C * A;
        }
        cout << ans << endl;
        return 0;
}

Submission Info

Judge Result