查看: 8464|回复: 5

红黑树python实现~

字体大小: 正常放大

表咯痴

19 主题	11 听众	63 积分

升级 61.05%

TA的每日心情

	慵懒 2015-5-5 10:06

签到天数: 12 天

[LV.3]偶尔看看II

自我介绍: 我就是我

电梯直达

1^#

发表于 2015-4-15 09:44 |只看该作者 |正序浏览

|招呼Ta 关注Ta

#!/usr/bin/env python
# -*- coding : utf-8 -*-

import os
import random

RED = 0
BLACK = 1

class Vector(object):
def __init__(self,x=None,y=None):
      self.x=x
      self.y=y

class Node(object):
      """docstring for Node"""
      def __init__(self,data=None,color=RED,left=None,right=None,parent=None):
         self.data = data
         self.color = color
         self.left = left
         self.right = right
         self.parent = parent

class RBtree(object):
def __init__(self):
      self.root=None
      self.size=0

def rb_rotate_left(self,node):
      right=node.right

      node.right=right.left
      if node.right is not None:
         right.left.parent=node

      right.left=node
      right.parent=node.parent

      if right.parent is None:

         self.root=right
      else:
         if node==node.parent.right:
               node.parent.right=right
         else:
               node.parent.left=right
      node.parent =right

def rb_rotate_right(self,node):
      left=node.left
      node.left=left.right

      if node.left is not None:
         left.right.parent=node

      left.right=node
      left.parent=node.parent

      if left.parent is None:
         self.root=left
      else:
         if node==node.parent.right:
            node.parent.right=left
         else:
            node.parent.left=left
      node.parent=left

def rb_insert_rebalance(self,node):
      parent=node.parent
      while parent and parent.color==RED:
         gparent = parent.parent
         if parent==gparent.left:
            uncle=gparent.right
            if uncle and uncle.color==RED:
                  uncle.color=BLACK
                  parent.color=BLACK
                  gparent.color=RED
                  node=gparent
            else:
                  if parent.right==node:
                     self.rb_rotate_left(parent)
                     tmp=parent
                     parent=node
                     node=tmp
                  parent.color=BLACK
                  gparent.color=RED

                  self.rb_rotate_right(gparent)

                  if uncle:
                     if uncle.right:
                        node=uncle.right
         else:
            uncle=gparent.left
            if uncle and uncle.color==RED:
                  uncle.color=BLACK
                  parent.color=BLACK
                  gparent.color=RED
                  node=gparent
            else:
                  if parent.left==node:
                     self.rb_rotate_right(parent)
                     tmp=parent
                     parent=node
                     node=tmp
                  parent.color=BLACK
                  gparent.color=RED
                  self.rb_rotate_left(gparent)

                  if uncle:
                     if uncle.left:
                        node=uncle.left
         parent=node.parent

      self.root.color=BLACK


def rb_search_auxiliary(self,node):
      tmp=self.root
      parent=None
      while tmp is not None:
         parent=tmp
         cmp=self.cmp(node,tmp)
         if cmp<0:
            tmp=tmp.left
         else:
            if cmp>0:
                  tmp=tmp.right
            else:
                  return tmp,parent

      return None,parent

def rb_insert(self,data):
      tmp=None
      node=Node(data)
      tmp,parent=self.rb_search_auxiliary(node)

      if tmp is not None:
         return

      node.parent =parent
      node.left=node.right=None
      node.color=RED

      if parent is not None:

         if self.cmp(parent,node)>0:
            parent.left=node
         else:
            parent.right=node
      else:
         self.root=node
      return self.rb_insert_rebalance(node)

def rb_erase_rebalance(self,node,parent):
      while((node is None or node.color==BLACK)and node !=self.root):
         if parent.left==node:
            other=parent.right
            if other.color==RED:
                  other.color=BALCK
                  parent.color=RED
                  self.rb_rotate_left(parent)
                  other=parent.right
            if (other.left is None or other.left.color==BLACK) and (other.right is None or other.right.color ==BLACK):
                  other.color=RED
                  node=parent
                  parent=node.parent
            else:
                  if other.right is None or other.right.color==BLACK:

                     if other.left is not None:
                        other.left.color=BLACK
                     other.color=RED
                     self.rb_rotate_right(other)
                     other=parent.right

                  other.color=parent.color
                  parent.color=BLACK
                  if other.right is not None:
                     other.right.color=BLACK
                  self.rb_rotate_left(parent)
                  node=self.root
                  break
         else:
            other=parent.left
            if other.color==RED:
                  other.color=BLACK
                  parent.color=RED
                  self.rb_rotate_right(parent)
                  other=parent.left
            if (other.left is None or other.left.color==BLACK) and (other.right is None or other.right.color==BLACK):
                  other.color=RED
                  node=parent
                  parent=node.parent
            else:
                  if other.left is None or other.left.color==BLACK:
                     if other.right is not None:
                        other.right.color=BLACK

                     other.color=RED
                     self.rb_rotate_left(other)
                     other=parent.left

                  other.color=parent.color
                  parent.color=BLACK

                  if other.left is not None:
                     other.left.color=BLACK

                  self.rb_rotate_right(parent)
                  node=self.root
                  break

      if node is not None:
         node.color=BLACK

def rb_erase(self,data):
      tmp_node=Node(data)
      node,parent = self.rb_search_auxiliary(tmp_node)
      if node is None:
         print "data is not exist."
         return

      old =node
      if node.left and node.right:
         node=node.right

         left=node.left
         while left is not None:
            node =left
            left=node.left

         child=node.right
         parent=node.parent
         color=node.color

         if child:
            child.parent=parent
         if parent:
            if parent.left==node:
                  parent.left=child
            else:
                  parent.right=child

         else:
            self.root=child

         if node.parent==old:
            parent=node
         node.parent=old.parent
         node.color=old.color
         node.right=old.right
         node.left=old.left

         if old.parent:
            if old.parent.left==old:
                  old.parent.left=node
            else:
                  old.parent.right=node
         else:
            self.root=node

         old.left.parent=node
         if old.right:
            old.right.parent=node

      else:
         if node.left is None:
            child =node.right
         else:
            if node.left is not None:
                  child=node.left
            else:
                  child=None
         parent=node.parent
         color=node.color
         if child:
            child.parent=parent
         if parent:
            if parent.left==node:
                  parent.left=child
            else:
                  parent.right=child
         else:
            self.root=child

      if color==BLACK:

         self.rb_erase_rebalance(child,parent)

def rb_travelse(self,node):
      if node is not None:
         print str(node.data)+'\t'+str(node.color)
         self.rb_travelse(node.left)
         if node.parent:
               if node.parent.color==0 and node.color==0:
                  print "error"
                  return
         self.rb_travelse(node.right)
         if node.parent:
               if node.parent.color==0 and node.color==0:
                  print "error"
                  return

      return


def cmp(self,node1,node2):
      if node1.data>node2.data :
         return 1
      if node1.data==node2.data :
         return 0
      if node1.data<node2.data:
         return -1

if __name__=="__main__":
print "main"
data=[28,6,39,78,6,43,61,56,71,38]
#for i in range(10):
# rand_num = random.randint(0, 100)
# data.append(rand_num)
#print data
t=RBtree()

for i in range(10):
      t.rb_insert(data[i])

t.rb_travelse(t.root)

print "---------------------------------"
t.rb_erase(data[7])

t.rb_travelse(t.root)

zan

python