c_cpp 实现disjoint_set数据结构

Posted 2021-05-14

#include <algorithm>
#include <iterator>
#include <unordered_set>
#include <vector>

class disjoint_set {
    struct ranked { int root, rank, size; };

public:
    explicit disjoint_set(int n): cc(n + 1) {
        for (int i = 0; i <= n; ++i) {
            cc[i].root = i;
            cc[i].rank = 0;
            cc[i].size = 1;
        }
    }

    std::size_t size(int x) const {
        return cc[x].size;
    }

    std::unordered_set<int> roots() {
        std::unordered_set<int> lst;

        std::transform(std::next(cc.begin()), cc.end(), std::inserter(lst, lst.begin()),
            [&](const ranked& r) {
                return find(r.root);
            });

        return lst;
    }

    void join(int x, int y) {
        ranked &xr = cc[find(x)];
        ranked &yr = cc[find(y)];

        if (xr.root == yr.root) {
            return;
        }

        if (xr.rank < yr.rank) {
            xr.root = yr.root;
            yr.size += xr.size;
        } else if (xr.rank > yr.rank) {
            yr.root = xr.root;
            xr.size += yr.size;
        } else {
            xr.root = yr.root;
            yr.rank++;
            yr.size += xr.size;
        }
    }

private:
    std::vector<ranked> cc;

    int find(int p) {
        if (cc[p].root == p) {
            return p;
        }
        cc[p].root = find(cc[p].root);
        return cc[p].root;
    }
};