Union-Find Problem

The problem is to find if a given edge (a, b) is connected in a graph with vertices {a, b, c, ...}.

Naive Solution – O(n^2)

• Maintain a parent array of size (len of vertex array)
• Update an edge (a, b) with the index a & b in the vertex array

ParentArray[a] = b
ParentArray[b] = b

Code

class Solution(object):
    def __init__(self, nodesCount):
        self.nodes = list(range(nodesCount + 1))

    def union(self, p, q):
        currentRoot = self.root(p)
        for index, num in enumerate(self.nodes):
            if num == currentRoot:
                self.nodes[index] = q
                
        print(self.nodes)

    def find(self, p, q):
        return self.root(p) == self.root(q)
    
    def root(self, p):
        return self.nodes[p]
    
if __name__ == "__main__":
    s = Solution(6)
    s.union(1, 2)
    s.union(2, 3)
    s.union(2, 4)
    
    print(s.find(1, 3))
    print(s.find(3, 4))
    print(s.find(3, 6))

Execution

$ python union-find.py 
[0, 2, 2, 3, 4, 5, 6]
[0, 3, 3, 3, 4, 5, 6]
[0, 4, 4, 4, 4, 5, 6]
True
True
False

Written with StackEdit.