Balance_Tree

# -*-coding: utf-8-*-

class TreeNode():
    def __init__(self, elem = None, lchild = None, rchild = None):
        self.elem = elem
        self.lchild = lchild
        self.rchild = rchild

class Tree():
    def __init__(self):
        self.root = None

    def add(self, elem = 'None'):
        new_node = TreeNode(elem)
        if self.root is None:
            self.root = new_node
        else:
            myQueue = [self.root]
            while True:
                cur = myQueue.pop(0)
                if cur.elem != 'None':
                    if cur.lchild is None:
                        cur.lchild = new_node
                        break
                    elif cur.rchild is None:
                        cur.rchild = new_node
                        break
                    else:
                        myQueue.append(cur.lchild)
                        myQueue.append(cur.rchild)


def printTree(root):
    print("Binary Tree:")
    printInOrder(root, 0, 'H', 17)
    
def printInOrder(root, height, s, length):
    if root is None:
        return
    printInOrder(root.rchild, height + 1, 'v', length)
    val = s + str(root.elem) + s
    lenM = len(val)
    lenL = (length - lenM) // 2
    lenR = length - lenM -lenL
    val = getSpace(lenL) + val + getSpace(lenR)
    print(getSpace(height * length) +val)
    printInOrder(root.lchild, height + 1, '^', length)

def getSpace(num):
    return ' ' * num

def isBalance(root):
    if root is None:
        return (True, 0)
    leftData = isBalance(root.lchild)
    if not leftData[0]:
        return (False, 0)
    rightData = isBalance(root.rchild)
    if not rightData[0]:
        return  (False, 0)
    if abs(rightData[1] - leftData[1]) > 1:
        return (False, 0)
    return (True, max(rightData[1], leftData[1]) + 1)
        
t1 = Tree()
l = [1, 'None', 2]
for each in l:
    t1.add(each)
printTree(t1.root)
print(isBalance(t1.root))