union_find.cpp

#include <vector>

// if parent[x] <  0, x is root and |parent[x]| is a number of the set
// if parent[x] >= 0, x is child
class UnionFind {
private:
    std::vector<int> parent;
    bool double_cnt;

public:
    UnionFind(std::size_t n, bool double_cnt_ = false)
        : parent(n, -1), double_cnt(double_cnt_) {}

    int find(int x) {
        if (parent[x] < 0)
            return x;
        else
            return (parent[x] = find(parent[x])); // path compression
    }

    void unite(int x, int y) {
        x = find(x);
        y = find(y);
        if (x == y && !double_cnt) return;
        if (parent[x] < parent[y]) { // |parent[x]| > |parent[y]|
            parent[x] += parent[y];
            if (x != y)
                parent[y] = x;
        } else {
            parent[y] += parent[x];
            if (x != y)
                parent[x] = y;
        }
    }
	
    int n_set(int x) {
        if (parent[x] >= 0) {
            parent[x] = find(parent[x]);
            x = parent[x];
        }
        return -parent[x]; // |parent[x]|
    }
};