leetcode-union-find

323. Number of Connected Components in an Undirected Graph.py

class UnionFindSet:
    def __init__(self, n):
        self._parents = [node for node in range(n)]
        self._ranks = [1 for _ in range(n)]
        
    def find(self, u):
        while u != self._parents[u]:
            self._parents[u] = self._parents[self._parents[u]]
            u = self._parents[u]
        return u
    
    def union(self, u, v):
        root_u, root_v = self.find(u), self.find(v)
        if root_u == root_v:
            return True
        
        if self._ranks[root_u] > self._ranks[root_v]:
            self._parents[root_v] = root_u
        elif self._ranks[root_v] > self._ranks[root_u]:
            self._parents[root_u] = root_v
        else:
            self._parents[root_u] = root_v
            self._ranks[root_v] += 1
        return False
    
class Solution:
    """
    Thought Process:
        Union-Find:
            step1: we iterate edge list. 
                For each egde, we mrege node a and node b together using union methord provided by Union-Find data structure.
            step2: we iterge all nodes.
                For each node, we get root node of this current node using find method provided by Union-Find data structure.
                At the same time, we use hashmap or hashset to store nb of root node.
            step3: number of uniue parent node we have is equal to nb of connected components we in this graph.
        Time complexity analysis:
            1. two pass. First pass is we go through m edges with union methods and second pass is we go through n nodes with find methods.
            2. For find and union methods, it can be thought as O(log*n) =~ O(1) amortized time.
            3. So, running time is O(m+n) in total.
    """
    def countComponents(self, n: int, edges: List[List[int]]) -> int:
        """ 
        T: O(m+n), where n is number of nodes and m is number of edges in the graph.
        S: O(n) for implementing Union-Find and hashmap
        """
        
        unionfind = UnionFindSet(n)
        # union
        for a, b in edges:
            unionfind.union(a, b)
        # find unique number of root node
        count = collections.defaultdict(int)
        for node in range(n):
            root_node = unionfind.find(node)
            count[root_node] += 1 # value represent nb of nodes in this component.
        return len(count)