平衡二叉树-python

#coding:utf8

# 全局变量，判断是否高度+1
taller = False
LH = 1
EH = 0
RH = -1

class BiTree (object):

	def __init__(self, data):
		self.data = data
		self.lChild = None
		self.rChild = None
		self.bf = 0

class AVLTree (object):

	def __init__(self):
		self.root = None

	def createAVL(self, inputList):
		if not inputList:
			return None
		for v in inputList:
			self.root = self.insertAVL(self.root, v)
		return self.root

	def leftRotate(self, parent):
		rchild = parent.rChild
		if not rchild:
			return
		parent.rChild = rchild.lChild
		rchild.lChild = parent
		return rchild

	def rightRotate(self, parent):
		lchild = parent.lChild
		if not lchild:
			return
		parent.lChild = lchild.rChild
		lchild.rChild = parent
		return lchild

	def leftBalance(self, parent):
		left = parent.lChild
		if LH == left.bf:
			parent.bf = EH
			left.bf = EH
			parent = self.rightRotate(parent)
		elif RH == left.bf:
			lr = left.rChild
			if EH == lr.bf:
				parent.bf = EH
				left.bf = EH
			elif LH == lr.bf:
				parent.bf = RH
				left.bf = EH
			elif RH == lr.bf:
				parent.bf = EH
				left.bf = LH
			lr.bf = EH
			parent.lChild = self.leftRotate(left)
			parent = self.rightRotate(parent)
		return parent

	def rightBalance(self, parent):
		right = parent.rChild
		if RH == right.bf:
			parent.bf = EH
			right.bf = EH
			parent = self.leftRotate(parent)
		elif LH == right.bf:
			rl = right.lChild
			if EH == rl.bf:
				parent.bf = EH
				right.bf = EH
			elif LH == rl.bf:
				parent.bf = EH
				right.bf = RH
			elif RH == rl.bf:
				parent.bf = LH
				right.bf = EH
			rl.bf = EH
			parent.rChild = self.rightRotate(right)
			parent = self.leftRotate(parent)
		return parent

	def insertAVL(self, parent, v):
		global taller
		if not parent:
			parent = BiTree(v)
			taller = True
			return parent
		if v < parent.data:
			parent.lChild = self.insertAVL(parent.lChild, v)
			if taller:
				if LH == parent.bf:
					parent = self.leftBalance(parent)
					taller = False
				elif EH == parent.bf:
					parent.bf += 1
					taller = True
				elif RH == parent.bf:
					parent.bf += 1
					taller = False
		else:
			parent.rChild= self.insertAVL(parent.rChild, v)
			if taller:
				if LH == parent.bf:
					parent.bf -= 1
					taller = False
				elif EH == parent.bf:
					parent.bf -= 1
					taller = True
				elif RH == parent.bf:
					parent = self.rightBalance(parent)
					taller = False
		return parent

	def midTraverse(self, root):
		if not root:
			return None
		self.midTraverse(root.lChild)
		print root.data
		self.midTraverse(root.rChild)

if __name__ == '__main__':
	T = AVLTree()
	l = [3, 2, 1, 4, 5, 6, 7, 10, 9, 8]
	root = T.createAVL(l)
	T.midTraverse(root)