Union Find Python template

union_find.py

"""
Union Find template for me.
"""
MAX_SIZE = 10000005

class UnionFind:
  
    def __init__(self):
        self.height = [0]*MAX_SIZE
        self.father=[0]*MAX_SIZE
        for i in range(MAX_SIZE):
            self.father[i]=i
    
    def find(self,a):
        if self.father[a] != a:
            self.father[a] = self.find(self.father[a])
        return self.father[a]
        
    def unite(self,a,b):
        a_val, b_val = self.find(a), self.find(b)
        if self.height[a_val]>self.height[b_val]:
            self.father[b_val]=a_val
            self.height[a_val] = max(self.height[a_val],self.height[b_val]+1)
        else:
            self.father[a_val]=b_val
            self.height[b_val]=max(self.height[b_val],self.height[a_val]+1)