Merging Communities

  • + 0 comments
    class UnionFind:
        def __init__(self, number_of_nodes):
            self.parent = [i for i in range(0, number_of_nodes+1)]
            self.rank = [1] * (number_of_nodes + 1)
            self.size = [1] * (number_of_nodes + 1)
        
        def find(self, node1):
            if self.parent[node1] == node1:
                return node1
            
            self.parent[node1] = self.find(self.parent[node1])
            return self.parent[node1]
                
        def union(self, node1, node2):
            root1, root2 = self.find(node1), self.find(node2)
            if root1 != root2:
                if self.rank[root1] < self.rank[root2]:
                    root1, root2 = root2, root1
                if self.rank[root1] == self.rank[root2]:
                    self.rank[root1] += 1
                self.parent[root2] = root1
                self.size[root1] = self.size[root1] + self.size[root2]
     
    inp_arr = input().split(' ')
    n, q = int(inp_arr[0]), int(inp_arr[1])
    
    uf = UnionFind(n)
    
    for i in range(q):
        inp_arr = input().split(' ')
        qtype = inp_arr[0]
        if qtype == 'Q':
            p = int(inp_arr[1])
            print(uf.size[uf.find(p)])
        else:
            p1, p2 = int(inp_arr[1]), int(inp_arr[2])
            uf.union(p1, p2)