1:#/usr/bin/env ruby
       2:class RedBlackTree
       3:# UPDATE 2006-01-30 see comments at end
       4:
Ruby   5:#Red/Black tree implementation based on 
       6:#Algorithms in C++, Sedgewick
       7:#Introduction To Algorithms Cormen, Thomas H. / Leiserson, Charles E . / Rivest, Ronald L . The MIT Press 07/1990
       8:# NOTE : LOOK AT END OF FILE TO SEE DIFFERENCES IN TRAVERSAL IDIOMS
       9:  RED   = 0
Ruby  10:  BLACK = 1
      11:
      12:# use 'protected' instead of 'private' so other instances can see these methods.
      13:# private REALLY means private in ruby
      14:protected
Ruby  15:
      16:  def copy(node)
      17:    @left  = node.left
      18:    @right = node.right
      19:    @val   = node.val
Ruby  20:    @color = node.color
      21:  end
      22:  
      23:  def rb_insert_left(val,sw)
      24:    if @left.nil? then 
Ruby  25:      @left = RedBlackTree.new(val)
      26:    else 
      27:      @left.rb_insert(val,sw)
      28:    end
      29:  end
Ruby  30:    
      31:  def rb_insert_right(val,sw)
      32:    if @right.nil? then
      33:      @right = RedBlackTree.new(val)
      34:    else
Ruby  35:      @right.rb_insert(val,sw)
      36:    end
      37:  end
      38:
      39:  def rot_left
Ruby  40:    if @right.nil? then return nil end
      41:    
      42:    # make the changes in a copy of current node __self
      43:    # on return, the caller will copy in 'me' to current node
      44:    me = RedBlackTree.new()
Ruby  45:    me.copy(self)
      46:    x        = me.right
      47:    me.right = x.left
      48:    x.left   = me
      49:    return   x
Ruby  50:  end
      51:
      52:  def rot_right
      53:    if @left.nil? then return nil end 
      54:       
Ruby  55:    # make the changes in a copy of current node __self
      56:    # on return, the caller will copy in 'me' to current node
      57:    me = RedBlackTree.new()
      58:    me.copy(self)
      59:    x       = me.left
Ruby  60:    me.left = x.right
      61:    x.right = me
      62:    return x
      63:  end
      64:       
Ruby  65:  def rb_insert(val,sw)
      66:    # if current node is a 4 node, split it by flipping its colors
      67:    # if both children of this node are red, change this one to red
      68:    # and the children to black
      69:    l = @left
Ruby  70:    r = @right
      71:    if (!l.nil? and(l.color==RedBlackTree::RED)and !r.nil? and(r.color==RedBlackTree::RED)) then
      72:      @color  = RedBlackTree::RED
      73:      l.color = RedBlackTree::BLACK
      74:      r.color = RedBlackTree::BLACK
Ruby  75:    end
      76:    
      77:    # go left or right depending on key relationship
      78:    if (val < @val) then
      79:      # recursively insert
Ruby  80:      rb_insert_left(val,0)
      81:      
      82:      # on the way back up check if a rotation is needed
      83:      # if search path has two red links with same orientation
      84:      # do a single rotation and flip the color bits
Ruby  85:      l = @left
      86:      if ((@color == RedBlackTree::RED)and !l.nil? and(l.color == RedBlackTree::RED)and(sw == 1)) then
      87:        x = rot_right()
      88:        copy(x) unless x.nil?
      89:      end
Ruby  90:       
      91:      # flip the color bits
      92:      l = @left
      93:      if !l.nil? then
      94:        ll = l.left
Ruby  95:        if !ll.nil? then
      96:          if ((l.color == RedBlackTree::RED)and(ll.color == RedBlackTree::RED)) then
      97:            x = rot_right()
      98:            copy(x) unless x.nil?
      99:            @color = RedBlackTree::BLACK
Ruby 100:            r = @right
     101:            r.color = RedBlackTree::RED unless r.nil?
     102:          end
     103:        end
     104:      end
Ruby 105:    else
     106:      rb_insert_right(val,1)
     107:      
     108:      # on the way back up check if a rotation is needed
     109:      # if search path has two red links with same orientation
Ruby 110:      # do a single rotation and flip the color bits
     111:      r = @right
     112:      if ((@color == RedBlackTree::RED)and !r.nil? and(r.color == RedBlackTree::RED)and(sw == 0)) then
     113:        x = rot_left()
     114:        copy(x) unless x.nil?
Ruby 115:      end
     116:       
     117:       # flip the color bits
     118:      r = @right
     119:      if !r.nil? then
Ruby 120:        rr = r.right
     121:        if !rr.nil? then
     122:          if ((r.color == RedBlackTree::RED)and(rr.color == RedBlackTree::RED)) then
     123:            x = rot_left()
     124:            copy(x) unless x.nil?
Ruby 125:            @color = RedBlackTree::BLACK
     126:            l = @left
     127:            l.color = RedBlackTree::RED unless l.nil?
     128:          end      
     129:        end
Ruby 130:      end
     131:    end
     132:  end
     133:  
     134:  # update 2006-01-30
Ruby 135:  # use attributes rather than writing explicit access functions
     136:  attr_reader :left,:right
     137:  attr_writer :left,:right,:color
     138:  
     139:# ============================================================
Ruby 140:# public methods
     141:# ============================================================
     142:public  
     143:  def initialize(val = nil)
     144:    @left    = nil
Ruby 145:    @right   = nil
     146:    @val     = val
     147:    @color   = RedBlackTree::RED
     148:  end
     149:  
Ruby 150:  def to_s()
     151:    return '[' + @val.to_s() + ',' + @color.to_s() + ']'
     152:  end    
     153:
     154:  # nice easy read-only accessors
Ruby 155:  attr_reader :val,:color
     156:  
     157:  def find(key)
     158:    result = nil
     159:    if (key == @val) then
Ruby 160:      result = self
     161:    elsif (key < @val) then
     162:      result = @left.find(key)  unless left.nil?
     163:    else
     164:      result = @right.find(key) unless right.nil?
Ruby 165:    end
     166:        return result
     167:  end
     168:       
     169:  # inorder traversal using a Ruby block as the visitor
Ruby 170:  # &v passes in the block as a Proc object but it is still available for 'yield' also
     171:  def inorder(depth,&v)
     172:    @left.inorder(depth+1,&v)  unless @left.nil?
     173:    yield self,depth
     174:    @right.inorder(depth+1,&v) unless @right.nil? 
Ruby 175:  end
     176:  
     177:  # main public insert
     178:  def insert(val)
     179:    rb_insert(val,0)
Ruby 180:    @color = RedBlackTree::BLACK
     181:  end
     182:  
     183:end
     184:
Ruby 185:
     186:# ==================================
     187:# test program
     188:# ==================================
     189:def testRedBlackTree
Ruby 190:  nodelist = [11,4,8,14,17,6,9,7,16,15,13,5,19,18,12,10,3,20,1]
     191:  root  = RedBlackTree.new(2)
     192:  nodelist.each {|n| root.insert(n)}
     193:  
     194:  # update 2006-01-30 removed lambda example because it is equivalent to the block example
Ruby 195:  # which is more idiomatic Ruby
     196:  
     197:  # use a block directly 
     198:  print "Ruby     = "  
     199:  root.inorder(0) {|n,d| print '(' + n.val().to_s() + ':' + n.color().to_s() + ':' + d.to_s() + '), '}
Ruby 200:  print "\n"
     201:  
     202:  # do a find
     203:  t = root.find(16)
     204:  print t
Ruby 205:  print "\n"
     206:end  
     207:
     208:testRedBlackTree
     209:
Ruby 210:# UPDATE 2006-01-30
     211:# fixed inorder traversal using blocks to be more conventional
     212:# used .nil? instead of == nil
     213:# reduced line count by using statement qualifiers
     214:#   I wonder if this makes the code less readable, but it seems to be the consensus
Ruby 215:#   from Ruby programmers who commented
     216:# changed indents to 2 spaces which is the Ruby convention
     217:#   again this looks less readable to my C/Java eyes
     218: