RedBlackTree.rb

#/usr/bin/env ruby

class RedBlackTree


#Red/Black tree implementation based on 
#Algorithms in C++, Sedgewick
#Introduction To Algorithms Cormen, Thomas H. / Leiserson, Charles E . / Rivest, Ronald L . The MIT Press 07/1990

  RED   = 0

  BLACK = 1



# use 'protected' instead of 'private' so other instances can see these methods.

# private REALLY means private in ruby

protected



  def copy(node)

    @left  = node.left

    @right = node.right

    @val   = node.val

    @color = node.color

  end

  

  def rb_insert_left(val,sw)

    if @left.nil? then 

      @left = RedBlackTree.new(val)

    else 

      @left.rb_insert(val,sw)

    end

  end

    

  def rb_insert_right(val,sw)

    if @right.nil? then

      @right = RedBlackTree.new(val)

    else

      @right.rb_insert(val,sw)

    end

  end



  def rot_left

    if @right.nil? then return nil end

    

    # make the changes in a copy of current node __self

    # on return, the caller will copy in 'me' to current node

    me = RedBlackTree.new()

    me.copy(self)

    x        = me.right

    me.right = x.left

    x.left   = me

    return   x

  end



  def rot_right

    if @left.nil? then return nil end 

       

    # make the changes in a copy of current node __self

    # on return, the caller will copy in 'me' to current node

    me = RedBlackTree.new()

    me.copy(self)

    x       = me.left

    me.left = x.right

    x.right = me

    return x

  end

       

  def rb_insert(val,sw)

    # if current node is a 4 node, split it by flipping its colors

    # if both children of this node are red, change this one to red

    # and the children to black

    l = @left

    r = @right

    if (!l.nil? and(l.color==RedBlackTree::RED)and !r.nil? and(r.color==RedBlackTree::RED)) then

      @color  = RedBlackTree::RED

      l.color = RedBlackTree::BLACK

      r.color = RedBlackTree::BLACK

    end

    

    # go left or right depending on key relationship

    if (val < @val) then

      # recursively insert

      rb_insert_left(val,0)

      

      # on the way back up check if a rotation is needed

      # if search path has two red links with same orientation

      # do a single rotation and flip the color bits

      l = @left

      if ((@color == RedBlackTree::RED)and !l.nil? and(l.color == RedBlackTree::RED)and(sw == 1)) then

        x = rot_right()

        copy(x) unless x.nil?

      end

       

      # flip the color bits

      l = @left

      if !l.nil? then

        ll = l.left

        if !ll.nil? then

          if ((l.color == RedBlackTree::RED)and(ll.color == RedBlackTree::RED)) then

            x = rot_right()

            copy(x) unless x.nil?

            @color = RedBlackTree::BLACK

            r = @right

            r.color = RedBlackTree::RED unless r.nil?

          end

        end

      end

    else

      rb_insert_right(val,1)

      

      # on the way back up check if a rotation is needed

      # if search path has two red links with same orientation

      # do a single rotation and flip the color bits

      r = @right

      if ((@color == RedBlackTree::RED)and !r.nil? and(r.color == RedBlackTree::RED)and(sw == 0)) then

        x = rot_left()

        copy(x) unless x.nil?

      end

       

       # flip the color bits

      r = @right

      if !r.nil? then

        rr = r.right

        if !rr.nil? then

          if ((r.color == RedBlackTree::RED)and(rr.color == RedBlackTree::RED)) then

            x = rot_left()

            copy(x) unless x.nil?

            @color = RedBlackTree::BLACK

            l = @left

            l.color = RedBlackTree::RED unless l.nil?

          end      

        end

      end

    end

  end

  

  # use attributes rather than writing explicit access functions

  attr_reader :left,:right

  attr_writer :left,:right,:color

  

# ============================================================

# public methods

# ============================================================

public  

  def initialize(val = nil)

    @left    = nil

    @right   = nil

    @val     = val

    @color   = RedBlackTree::RED

  end

  

  def to_s()

    return '[' + @val.to_s() + ',' + @color.to_s() + ']'

  end    



  # nice easy read-only accessors

  attr_reader :val,:color

  

  def find(key)

    result = nil

    if (key == @val) then

      result = self

    elsif (key < @val) then

      result = @left.find(key)  unless left.nil?

    else

      result = @right.find(key) unless right.nil?

    end

		return result

  end

       

  # inorder traversal using a Ruby block as the visitor

  # &v passes in the block as a Proc object but it is still available for 'yield' also

  def inorder(depth,&v)

    @left.inorder(depth+1,&v)  unless @left.nil?

    yield self,depth

    @right.inorder(depth+1,&v) unless @right.nil? 

  end

  

  # main public insert

  def insert(val)

    rb_insert(val,0)

    @color = RedBlackTree::BLACK

  end

  

end



# ==================================

# test program

# ==================================

def testRedBlackTree

  nodelist = [11,4,8,14,17,6,9,7,16,15,13,5,19,18,12,10,3,20,1]

  root  = RedBlackTree.new(2)

  nodelist.each {|n| root.insert(n)}

  

  # removed lambda example because it is equivalent to the block example

  # which is more idiomatic Ruby

  

  # use a block directly 

  print "Ruby     = "  

  root.inorder(0) {|n,d| print '(' + n.val().to_s() + ':' + n.color().to_s() + ':' + d.to_s() + '), '}

  print "\n"

  

  # do a find

  t = root.find(16)

  print t

  print "\n"

end  



testRedBlackTree