Union Find

這個資料結構可以用來讓 Minimum spanning tree 的求取過程加速,一個簡單的實作如下:

#include <iostream>
#include <vector>

using namespace std;

class UnionFind {
public:
    UnionFind(int n) {
        parent.reserve(n);
        for(int i = 0; i < n; i++) {
            parent[i] = i;
        }
    }

    void connect(int nodeA, int nodeB) {
        int rootA = find(nodeA);
        int rootB = find(nodeB);

        if(rootA != rootB) {
            parent[nodeB] = rootA;
        }
    };

    int find(int node) {
        if(parent[node] == node)
            return node;

        vector<int> path;
        while(node != parent[node]) {
            path.push_back(node);
            node = parent[node];
        }

        for(auto n: path) {
            parent[n] = node;
        }

        return node;
    };

    bool query(int nodeA, int nodeB) {
        int rootA = find(nodeA);
        int rootB = find(nodeB);
        return rootA == rootB;
    }

private:
    vector<int> parent;
};

int main(int argc, const char * argv[]) {
    UnionFind uf(10);
    cout << uf.query(1, 2) << endl; // false

    uf.connect(1, 2);
    cout << uf.query(1, 2) << endl; // true
    cout << uf.query(1, 3) << endl; // false

    uf.connect(2, 5);
    cout << uf.find(5) << endl; // 1

    return 0;
}

Last updated