Clean union find python code

  • 2
    class UnionFind(object):
        def __init__(self, n):
            self.count = n
            self.size = [1] * n
   = range(n)
        def find(self, p):
            while[p] != p:
      [p] =[[p]]
                p =[p]
            return p
        def union(self, p, q):
            idp, idq = map(self.find, (p, q))
            if idp == idq:
            less, more = (
                (idp, idq) if self.size[idp] < self.size[idq] else (idq, idp))
  [less] =[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.