WeightedQuickUnionUF.py

class WeightedQuickUnionUF(object):
    def __init__(self,N):
        self.__id = []
        self.sz = []
        for i in range(N):
            self.__id.append(i)
            self.sz.append(1)
            
    def root(self,i):
        while(i!=self.__id[i]):
            self.__id[i] = self.__id[self.__id[i]]
            i = self.__id[i]
        return i
        
    def connected(self,p,q):
        return self.root(p) == self.root(q)
        
    def union(self,p,q):
        i = self.root(p)
        j = self.root(q)
        if (self.sz[i] > self.sz[j]):
            self.__id[i] = j
            self.sz[j] = self.sz[j] + self.sz[i]
        else:
            self.__id[j] = i
            self.sz[i] = self.sz[i] + self.sz[j]
            
        
if __name__ == "__main__":
     u = WeightedQuickUnionUF(4)
     u.union(0,1)
     u.union(0,2)
     print "0 and 3 is connect? %s" % u.connected(0,3)