平衡二叉树

#!/usr/bin/env python
#coding=utf-8
import sys


#-------------------------------------------------------------------------------
#--------------------------- class BinaryTree define #---------------------------
#-------------------------------------------------------------------------------
class BTree(object):
   """
   base binary tree class
   three traverse function and get height function
   """
   def __init__(self,v,l,r,p):
       self.value = v
       self.lchild = l
       self.rchild = r
       self.parent = p
       self.bf = 0

   def show(self):
       """
       print the BinaryTree
       """
       print self.value
   
   def PreTrave(self):
       """
       previous order traverse the binary tree
       """
       if self.value is not None:
           self.show()
           if self.lchild is not None:
               self.lchild.PreTrave()
           if self.rchild is not None:
               self.rchild.PreTrave()
           
   def MidTrave(self):
       """
       Middle order traverse the binary tree
       """
       if self.value is not None:
           if self.lchild is not None:
               self.lchild.MidTrave()
           self.show()
           if self.rchild is not None:
               self.rchild.MidTrave()

   def PosTrave(self):
       """
       Post order traverse the binary tree
       """     
       if self.value is not None:
           if self.lchild is not None:
               self.lchild.PosTrave()
           if self.rchild is not None:
               self.rchild.PosTrave()
           self.show()
       
   def GetHeight(self):
       """
       get the height of the binarytree
       """
       if self.lchild is not None:
           hl = self.lchild.GetHeight()
       else:
           hl = 0
       if self.rchild is not None:
           hr = self.rchild.GetHeight()
       else:
           hr = 0
       return ((hl > hr) and hl or hr) + 1
       # return the larger height of child plus one

           


#-------------------------------------------------------------------------------
#---------------------------- Binary Search Tree #-------------------------------
#-------------------------------------------------------------------------------
class BSTree(BTree):
   """
   binary search tree , could add , delete but not balanced
   derived from class BTree(base binary tree)
   """
   def Find(self,key):
       global ptr
       global current
       ptr = self
       while ptr:
           if ptr.value == key:
               break
           if ptr.value < key:
               current = ptr
               ptr = ptr.rchild
           else:
               current = ptr
               ptr = ptr.lchild
       return ptr
   

   def Insert(self,key):
       global ptr
       global current
       find = self.Find(key)
       if find:
           return False
       global mode
       if mode == 2:
           ptr = BSTree(key,None,None,current)
       if mode == 1:
           ptr = AVLTree(key,None,None,current)
       if current.value > key:
           current.lchild = ptr
           ptr.parent = current
       else:
           current.rchild = ptr
           ptr.parent = current
       return True


   def RmvNode(self,key):
       if self.value == key:
           root = None
           self = None
           return False
       return self.__RmvNode(self.Find(key))


   def __RmvNode(self,ptr):
       global current
       if self.lchild is None or self.rchild is None:
           if self == ptr:
               self = (self.lchild is not None) and self.lchild or self.rchild
               current = self
               current.parent = None
               self.parent = None
               return True
       if ptr is None:
           return False
       tmp = ptr
       if ptr.lchild is None or ptr.rchild is None:
           if ptr.parent is None:
               root = ptr.lchild is None and ptr.rchild or ptr.lchild
               return True
           if ptr.lchild is None:
               if ptr.parent.lchild == ptr:
                   ptr.parent.lchild = ptr.rchild
               else:
                   ptr.parent.rchild = ptr.rchild
               ptr = ptr.rchild
           else:
               if ptr.parent.lchild == ptr:
                   ptr.parent.lchild = ptr.lchild
               else:
                   ptr.parent.rchild = ptr.lchild
               ptr = ptr.lchild
               if ptr:
                   ptr.parent = current
           return True
       if ptr.bf < 0:
           ptr = ptr.rchild
           while ptr.lchild:
               ptr = ptr.lchild
       else:
           ptr = ptr.lchild
           while ptr.rchild:
               ptr = ptr.rchild
       tmp.value = ptr.value
       current = ptr.parent
       if current.lchild == ptr:
           current.lchild = None
       else:
           current.rchild = None
       return True
       

#-------------------------------------------------------------------------------
#---------------------------------- AVL Tree #-----------------------------------
#-------------------------------------------------------------------------------
class AVLTree(BSTree):
   """
   AVL tree derived from BSTree(binary search tree)
   balanced dynamicly when data added and deleted
   """
   def AVLInsert(self,key):
       if self.Insert(key) == False:
           return False
       global ptr
       global current
       ptr = key
       while current is not None:
           cp = current.parent
           croot = (cp is not None) and ((cp.lchild == current) and (cp.lchild) or (cp.rchild)) or (root) 
           current.bf += (current.value > ptr) and 1 or -1
           if current.bf == -2:
               croot = croot.L_Balance()
           elif current.bf == 2:
               croot = croot.R_Balance()
           if current.bf == 0:
               break
           ptr = current.value
           current = current.parent
       while self.parent:
           self = self.parent
       return True
   
   
   def Remove(self,key):
       global current
       global rmvresult
       if self.RmvNode(key) == False:
           print "the tree is empty or no such node"
           return False
       while current:
           hl = (current.lchild) and current.lchild.GetHeight() or 0
           hr = (current.rchild) and current.rchild.GetHeight() or 0
           current.bf = hl - hr
           if current.bf == -2:
               current = current.L_Balance()
           elif current.bf == 2:
               current = current.R_Balance()
           if current.parent is None:
               break
           ptr = current
           current = current.parent
       self = current
       rmvresult = self
       return True


   def L_Balance(self):
       global current
       if (isinstance(self.rchild,AVLTree)) and (self.rchild.bf) == 1:
           self.rchild = (self.rchild).R_Rotate()
       self = self.L_Rotate()
       current = self
       return self

   def R_Balance(self):
       global current
       if (isinstance(self.lchild,AVLTree))and(self.lchild.bf) == -1:
           self.lchild = (self.lchild).L_Rotate()
       self = self.R_Rotate() 
       current = self
       return self


   def L_Rotate(self):
       tmp = self.rchild
       tmp.parent = (self.parent is not None) and (self.parent) or None
       if self.parent and self.parent.lchild == self:
           self.parent.lchild = tmp
       if self.parent and self.parent.rchild == self:
           self.parent.rchild = tmp
       self.parent = tmp
       self.rchild = (tmp.lchild is not None) and (tmp.lchild) or None
       if tmp.lchild:
           tmp.lchild.parent = self
       tmp.lchild = self
       self.bf += 1
       self.bf -= (tmp.bf < 0) and tmp.bf or 0
       tmp.bf += 1
       tmp.bf += (self.bf > 0) and self.bf or 0
       return tmp
  
  
   def R_Rotate(self):
       tmp = self.lchild
       tmp.parent = (self.parent is not None) and (self.parent) or None
       if self.parent and self.parent.lchild == self:
           self.parent.lchild = tmp
       if self.parent and self.parent.rchild == self:
           self.parent.rchild = tmp
       self.parent = tmp
       self.lchild = (tmp.rchild is not None) and tmp.rchild or None
       if tmp.rchild:
           tmp.rchild.parent = self
       tmp.rchild = self
       self.bf -= 1
       self.bf -= (tmp.bf > 0) and tmp.bf or 0
       tmp.bf -= 1
       tmp.bf += (self.bf < 0) and self.bf or 0
       return tmp
       


#-------------------------------------------------------------------------------
#------------------------------------ test #-------------------------------------
#-------------------------------------------------------------------------------
if __name__ == '__main__':
  print "1:avl tree , 2:binary search tree"
  while 1:
       try:
           t = raw_input()
           if int(t) == 1:
               mode = 1
               break
           if int(t) == 2:
               mode = 2
               break
       except ValueError:
           print "please choose 1 or 2"
  print "insert a number"
  t = raw_input()
  while t:
       try:

           if mode == 1:

               root = AVLTree((int(t)),None,None,None)
               break

           if mode == 2:

               root = BSTree((int(t)),None,None,None)
               break
       except ValueError:
           print "only int value acceptable"
  #1:avl tree , 2: bs tree
  while 1:
       try:
           print "insert a number"
           t = raw_input()
           if t == 'quit':
               break
           if mode == 1:
               if root.AVLInsert(int(t)) == False:
                   print "already exists"
               else:
                   while root.parent:
                       root = root.parent
                   print "now the avl tree:"
                   root.MidTrave()
           if mode == 2:
               if root.Insert(int(t)) == False:
                   print "already exists"
               else:
                   while root.parent:
                       root = root.parent
                   print "now the binary search tree:"
                   root.MidTrave()
       except ValueError:
           print "only int value acceptable"
  print "mid"
  root.MidTrave()
  print "now test delete function"
  while 1:
       print "which number you want to delete"
       t = raw_input()
       if t == 'quit':
           break
           try:
               if mode == 1:
                   if root.Remove(int(t)) == False:
                       print "no such node"
                   else:
                       root = rmvresult
               if mode == 2 and root.RmvNode(int(t)) == False:
                   print "no such node"
               print "after delete:"
               root.MidTrave()
           except ValueError:
           print "only int value acceptable"

来公司写的一个东西，当时是为了学习python
没想到还可以现在用来交作业。。。
说不定将来会用上