Union find is a algorithm that widly used in check connectivity in a graph.

Imagine connecting node as tree. Need to balance it while union and find, or else the time complexity for traversal would be O(N) instead of the height of the tree.

Sample Implementation

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
class UF:
    def __init__(self, n=0):
        # n disconnected node
        # number of tree
        self.count = n

        # the parent of each node
        self.parent = {}

        # size of each tree
        self.size = {}

        for i in range(n):
            self.parent[i] = i
            self.size[i] = 1

    def union(self, p, q):
        """
        connect p and q
        Parameter:
            p: int
            q: int
        """
        rootP = self.find(p)
        rootQ = self.find(q)
        if rootP == rootQ:
            return

        # for balance purpose 
        # append the smaller party to the larger party
        if self.size[rootP] > self.size[rootQ]:
            self.parent[rootQ] = rootP
            self.size[rootP] += self.size[rootQ]
        else:
            self.parent[rootP] = rootQ
            self.size[rootQ] += self.size[rootQ]
        self.count -= 1

    def connected(self, p, q):
        """
        check connectivity between p and q
        """
        return self.find(p) == self.find(q)

    def find(self, x):
        """
        find the root of x
        """
        while self.parent[x] != x:

            # balance the tree while finding
            self.parent[x] = self.parent[self.parent[x]]
            x = self.parent[x]
        return x

Reference

wiki
labuladong - union find (in Chinese)