实现一个并查集结构

#include <iostream>
#include <vector>

using namespace std;

class DisjointSetUnion {
private:
    vector<int> parent;
    vector<int> rank;

public:
    DisjointSetUnion(int n) {
        parent.resize(n);
        rank.resize(n, 0);
        for (int i = 0; i < n; ++i) {
            parent[i] = i;
        }
    }

    int find(int x) {
        if (parent[x] != x) {
            parent[x] = find(parent[x]);
        }
        return parent[x];
    }

    void unite(int x, int y) {
        int rootX = find(x);
        int rootY = find(y);
        if (rootX == rootY) return;
        
        if (rank[rootX] < rank[rootY]) {
            parent[rootX] = rootY;
        } else if (rank[rootX] > rank[rootY]) {
            parent[rootY] = rootX;
        } else {
            parent[rootY] = rootX;
            rank[rootX]++;
        }
    }

    bool same(int x, int y) {
        return find(x) == find(y);
    }
};

int main() {
    int n = 5; // 设置并查集的大小
    DisjointSetUnion dsu(n);

    // 进行一系列的合并操作
    dsu.unite(0, 1);
    dsu.unite(1, 2);
    dsu.unite(3, 4);

    // 输出每个元素所在集合的代表元素
    for (int i = 0; i < n; ++i) {
        cout << "Element " << i << " belongs to set with representative " << dsu.find(i) << endl;
    }

    // 检查元素是否在同一个集合中
    cout << "Are 0 and 2 in the same set? " << (dsu.same(0, 2) ? "Yes" : "No") << endl;
    cout << "Are 1 and 4 in the same set? " << (dsu.same(1, 4) ? "Yes" : "No") << endl;

    return 0;
}