Clean union find python code


  • 2
    class UnionFind(object):
        def __init__(self, n):
            self.count = n
            self.size = [1] * n
            self.id = range(n)
    
        def find(self, p):
            while self.id[p] != p:
                self.id[p] = self.id[self.id[p]]
                p = self.id[p]
            return p
    
        def union(self, p, q):
            idp, idq = map(self.find, (p, q))
            if idp == idq:
                return
    
            less, more = (
                (idp, idq) if self.size[idp] < self.size[idq] else (idq, idp))
    
            self.id[less] = self.id[more]
            self.size[more] += self.size[less]
    
            self.count -= 1
    
    
    class Solution(object):
        def countComponents(self, n, edges):
            unionFind = UnionFind(n)
            [unionFind.union(*e) for e in edges]
            return unionFind.count

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.