Disjoint Set / Union Find Template

tags: algorithm template c++ Union Find

C++

class UnionFind {
public:
    UnionFind(int size) : root(size), rank(size) {
        for (int i = 0; i < size; i++) {
            root[i] = i;
            rank[i] = 1;
        }
    }

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

    void unionSet(int x, int y) {
        int rootX = find(x);
        int rootY = find(y);
        if (rootX != rootY) {
            if (rank[rootX] > rank[rootY]) {
                root[rootY] = root[rootX];
            } else if (rank[rootX] < rank[rootY]) {
                root[rootX] = root[rootY];
            } else {
                root[rootY] = root[rootX];
                rank[rootX] += 1;
            }
        }
    }

    bool isConnected(int x, int y) {
        return find(x) == find(y);
    }
private:
    vector<int> root;
    vector<int> rank;
};

Python

class UnionFind:
    def __init__(self, size):
        self.root = [i for i in range(size)]
        self.rank = [1] * size

    def find(self, x):
        if x == self.root[x]:
            return x
        self.root[x] = self.find(self.root[x])
        return self.root[x]

    def union(self, x, y):
        rootX = self.find(x)
        rootY = self.find(y)
        if rootX != rootY:
            if self.rank[rootX] > self.rank[rootY]:
                self.root[rootY] = rootX
            elif self.rank[rootX] < self.rank[rootY]:
                self.root[rootX] = rootY
            else:
                self.root[rootY] = rootX
                self.rank[rootX] += 1

    def is_connected(self, x, y):
        return self.find(x) == self.find(y)

Related Topics

LC 1971. Find if Path Exists in Graph

LC 721. Accounts Merge

Complexity

n = the number of vertices in the graph.

	Time Complexity	Space Complexity
Constructor	O(n)	O(n)
Find	O(α(n))	O(1)
Union	O(α(n))	O(1)
isConnected	O(α(n))	O(1)

Note:

• n is the number of vertices in the graph. α refers to the Inverse Ackermann function. In practice, we assume it's a constant. In other words, O(α(n)) is regarded as O(1) as average.
• When using the combination of union by rank and the path compression optimization, the find operation will take O(α(n)) time on average. Since union and isConnected both make calls to find and all other operations require constant time, union and isConnected functions will also take O(α(n)) time on average.

Note

Leetcode
代碼隨想錄 - 并查集理论基础