Programming Practice: Alliances in Hogwarts

p13105.cpp

#include <iostream>
#include <vector>

class DisjointSet {

private:
    std::vector<int> parent, rank;

public:
    DisjointSet(int n) : parent(2 * n, -1), rank(2 * n, -1) {
    }

    int find(int x) {
        // unassigned parent
        if (parent[x] == -1) {
            return -1;
        }
        if (parent[x] == x) {
            // is root
            return x;
        } else {
            // find root
            parent[x] = find(parent[x]);
            return parent[x];
        }
    }

    void join(int x, int y) {
        int xroot = find(x);
        int yroot = find(y);
        // already joined
        if (xroot == yroot && xroot != -1) {
            return;
        }
        // merge xroot
        if (xroot == -1) {
            xroot = x;
            parent[x] = x;
        }
        // merge yroot
        if (yroot == -1) {
            yroot = y;
            parent[y] = y;
        }
        // merge into higher rank
        if (rank[xroot] > rank[yroot]) {
            rank[xroot] += rank[yroot];
            parent[yroot] = xroot;
        } else {
            rank[yroot] += rank[xroot];
            parent[xroot] = yroot;
        }
    }

    void makeFriend(int x, int y) {
        // join friend and enemy sets
        join(x * 2, y * 2);
        join(x * 2 + 1, y * 2 + 1);
    }

    void makeEnemy(int x, int y) {
        // enemy of my enemy is my friend
        join(x * 2, y * 2 + 1);
        join(x * 2 + 1, y * 2);
    }

    int isEnemy(int x, int y) {
        int xp = find(x * 2);
        int yp = find(y * 2);
        int yn = find(y * 2 + 1);
        if (xp == -1 || yp == -1)
            // undefined
            return -1;
        if (xp == yn)
            // enemies
            return 1;
        return 0;
    };

    int isFriend(int x, int y) {
        int xp = find(x * 2);
        int yp = find(y * 2);
        if (x == y)
            // friend of self
            return 1;
        if (xp == -1 || yp == -1)
            // undefined
            return -1;
        if (xp == yp)
            // friends
            return 1;
        return 0;
    }
};


int main() {
    int n;
    while(std::cin >> n) {
        int q;
        std::cin >> q;
        DisjointSet d(n);
        int cmd, x, y;
        for (int i = 0; i < q; i++) {
            std::cin >> cmd >> x >> y;
            switch (cmd) {
                case 1:
                    if (d.isEnemy(x, y) == 1) {
                        std::cout << -1 << std::endl;
                    } else {
                        d.makeFriend(x, y);
                    }
                    break;
                case 2:
                    if (d.isFriend(x, y) == 1) {
                        std::cout << -1 << std::endl;
                    } else {
                        d.makeEnemy(x, y);
                    }
                    break;
                case 3:
                    std::cout << (d.isFriend(x, y) > 0 ? 1 : 0) << std::endl;
                    break;
                default:
                    std::cout << (d.isEnemy(x, y) > 0 ? 1 : 0) << std::endl;
                    break;
            }
        }
    }

    return 0;
}