Weighted quick-union algorithm : This is an improvement of quick union algorithm. In this we keep track of size of a tree and when we merge teo trees we always make the root of smaller tree to point at the root of bigger tree. complexity : O(lg N)

weighed_quick_union.py

"""
This is an improvement over quick union algorithm. In this we keep track of size
of a tree and when we merge two trees we always make the root of smaller tree
to point at the root of bigger tree.

Link : https://class.coursera.org/algs4partI-002/lecture/8
"""

tree_sz=[1]*10           #initialize a size array

def get_root(lis,item):
    """
    return the root of an item
    """
    root=item
    while lis[root] != root:
        root = lis[root]
    return root    
def connected(lis,item1,item2):
    """
    Checks if two items are connected
    or not.
    Connected means both items have the
    same root
    """
    root1=get_root(lis,item1)
    root2= get_root(lis,item2)
    return True if root1==root2 else False

def quick_union(lis,item1,item2):
    """
    Join two nodes(join the root of smaller tree to root of bigger tree)
    """
    root1=get_root(lis,item1)
    root2= get_root(lis,item2)
    
    if root1 != root2:
        if tree_sz[root1]==tree_sz[root2]:       #if size of both tree is same
            lis[root2]=root1
            tree_sz[root1]+=tree_sz[root2]
            tree_sz[root2]=0                      # set the size of tree was merged to 0

        elif tree_sz[root1] < tree_sz[root2]:   
            lis[root1]=root2
            tree_sz[root2]+=tree_sz[root1]
            tree_sz[root1]=0

        elif tree_sz[root1] > tree_sz[root2]:
            lis[root2]=root1
            tree_sz[root1]+=tree_sz[root2]
            tree_sz[root2]=0    
    else:
        print "Already connected"