rbTree

rbTree.py

class Node( object ):
  """
  .. class:: Node
  
  rbTree Node
  """

  def __init__( self, key ):
    """  c'tor """
    self.left = None 
    self.right = None
    self.red = False
    self.key = key
    self._parent = None

  @property
  def parent(self):
    return self._parent 

  @parent.setter
  def parent( self, parent ):
    self._parent = parent 
  
  def grandpa( self ):
    """ grandparent """
    if self.parent:
      return self.parent.parent

  def uncle( self ):
    """ uncle """
    grandpa = self.grandpa()
    if not grandpa:
      return None
    if self.parent == grandpa.left:
      return grandpa.right
    return grandpa.left
 

class rbTree( object ):
  """
  .. class:: rbTree

  as seen in T. Cormen, C. Leiseron, R. Rivest, C. Stein 
  """

  def __init__( self ):
    """ c'tor """
    # # dummy node 
    self._null = Node(None)
    self._null.left = self._null
    self._null.right = self._null
    self._null.parent = self._null
    # # initial tree root
    self.root = self._null

  def insert( self, key ):
    """ insert key :key: into the tree """

    y = self._null
    x = self.root
    while x != self._null:
      y = x
      if key <= x.key:
        x = x.left
      else:
        x = x.right
    
    node = Node( key )
    node.parent, node.left, node.right, node.red  = y, self._null, self._null, True 
    
    if y == self._null:
      self.root = node
    elif node.key <= y.key:
      y.left = node
    else:
      y.right = node

    self._fixIns( node )
      
  def min( self, node = None ):
    """ find min from a given node """
    node = node if node else self.root
    while node.left.key:
      node = node.left
    return node.key
  
  def max( self, node = None ):
    """ find max from a given node """
    node = node if node else self.root
    while node.right.key:
      node = node.right
    return node.key
      
  def _fixIns( self, node ):
    """ fixtures after insertion """

    while node.parent.red:
      
      if node.parent == node.parent.parent.left:
      
        if node.parent.parent.right.red:

          node.parent.red = False
          node.parent.parent.right.red = False
          node.parent.parent.red = True 
          node = node.parent.parent
        else:
          if node == node.parent.right:
            node = node.parent
            self.rol(node)
          node.parent.red = False
          node.parent.parent.red = True
          self.ror( node.parent.parent )

      else:
      
        if node.parent.parent.left.red:

          node.parent.red = False
          node.parent.parent.left.red = False
          node.parent.parent.red = True
          node = node.parent.parent

        else:

          if node == node.parent.left:
            node = node.parent
            self.ror(node)

          node.parent.red = False
          node.parent.parent.red = True
          self.rol( node.parent.parent )

    self.root.red = False
        
  def rol( self, x ):
    """ rotate left """
    if not x.key: return
    y = x.right
    x.right = y.left
    if y.left != self._null:
      y.left.parent = x
    y.parent = x.parent
    if x.parent == self._null:
      self.root = y
    elif x == x.parent.left: 
      x.parent.left = y
    else:
      x.parent.right = y
    y.left = x
    x.parent = y

  def ror( self, y ):
    """ rotate right """
    if not y.key: return
    x = y.left
    y.left = x.right
    if x.right != self._null:
      x.right.parent = y
    x.parent = y.parent
    if y.parent == self._null:
      self.root = x
    elif y == y.parent.right: 
      y.parent.right = x
    else:
      y.parent.left = x
    x.right = y
    y.parent = x

  def transplant( self, u, v ):
    """ transplant subtrees 
    plant v in place of u 
    """
    if u.parent == self._null:
      self.root = v
    elif u == u.parent.left:
      u.parent.left = v
    else:
      u.parent.right = v
    v.parent = u.parent

  def find( self, k ):
    """ find for a key k """
    node = self.root
    while node != self._null and node.key != k:
      if k < node.key:
        node = node.left
      else:
        node = node.right
    return node

  def delete( self, z ):
    """ delete Node z """
    y = z
    zcol = y.red
    if z.left == self._null:
      x = z.right
      self.transplant( z, z.right )
    elif z.right == self._null:
      x = z.left
      self.transplant( z, z.left )
    elif y == self.min( z.right ):
      zcol = y.red
      x = y.right
      if y.parent == z:
        x.parent = y
      else:
        self.transplant( y, y.right )
        y.right = z.right
        y.right.parent = y
      self.transplant( z, y )
      y.left = z.left
      y.left.parent = y
      y.red = z.red
    if not zcol:
      self._fixDel( x )
      
  def _fixDel( self, x ):
    """ fixtures after deletion """
    while x != self.root and x.red:
      if x == x.parent.left:
        w = x.parent.right
        if w.red:
          w.red = False
          x.parent.red = True
          self.rol( x.parent )
          w = x.parent.right
        if not w.left.red and not w.right.red:
          w.red = True
          x = x.parent
        else:
          if not w.right.red:
            w.left.red = False
            w.red = True
            self.ror(w)
            w = x.parent.right
          w.red = x.parent.red
          x.parent.red = False
          w.right.red = False
          self.rol( x.parent )
          x = self.root
      else:
        w = x.parent.left
        if w.red:
          w.red = False
          x.parent.red = True
          self.ror( x.parent )
          w = x.parent.left
        if not w.left.red and not w.right.red:
          w.red = True
          x = x.parent
        else:
          if not w.left.red:
            w.right.red = False
            w.red = True
            self.rol(w)
            w = x.parent.left
          w.red = x.parent.red
          x.parent.red = False
          w.left.red = False
          self.ror( x.parent )
          x = self.root
      
  def dfs( self, node = None ):
    """ dummy inorder dfs to check """
    node = node if node else self.root
    if node.left.key:
      self.dfs( node.left )
    print node.key, 
    if node.right.key:
      self.dfs( node.right )
    
  def dot( self, fname ):
    """ dummy dot file creation 

    to generate pic run:
      neato -T png  fname > fname.png
    or 
      dot -T png fname > fname.png
    """
    f = open( fname, "w+" )
    f.write( "digraph rb {\n")
    def nodeId( node ):
      return "node%s" % id(node)
    def nodeCol( node ):
      return "red" if node.red else "black"
    def visit( node, f ):
      f.write( '%s [label="%s",color="%s",shape="circle"];\n' % ( nodeId(node), node.key, nodeCol(node) ) )
      if node.left.key:
        visit( node.left, f )
        f.write( "%s -> %s;\n" % ( nodeId(node), nodeId(node.left) ) )
      if node.right.key:
        visit( node.right, f )
        f.write( "%s -> %s;\n" % ( nodeId(node), nodeId(node.right) ) )
    visit( self.root, f )            
    f.write( "}\n" )
    f.close()

# # just for fun make some dot 
if __name__ == "__main__":

  t = rbTree()
  x = range(500)
  import random
  random.shuffle( x )
  for i in x:
    t.insert( i )
  t.dot( "rb500.dot")