#!/usr/bin/env ruby
class RedBlackTree
  include Enumerable
#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
# NOTE : LOOK AT END OF FILE TO SEE DIFFERENCES IN TRAVERSAL IDIOMS
  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?
      @left = RedBlackTree.new(val)
    else 
      @left.rb_insert(val,sw)
    end
  end
    
  def rb_insert_right(val,sw)
    if @right.nil?
      @right = RedBlackTree.new(val)
    else
      @right.rb_insert(val,sw)
    end
  end

  def rot_left
    return nil if @right.nil?
    # make the changes in a copy of current node __self
    # on return, the caller will copy in 'me' to current node
    me = self.dup
    x          = me.right
    me.right   = x.left
    x.left     = me
    return   x
  end

  def rot_right
    return nil if @left.nil?       
    # make the changes in a copy of current node __self
    # on return, the caller will copy in 'me' to current node
    me = self.dup
    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
    [self, @left, @right].each { |t| t.flip } if !@left.nil? and @left.red? and !@right.nil? and @right.red?
    
    # go left or right depending on key relationship
    if val < @val
      # 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
      copy(rot_right() || self) if red? and !@left.nil? and @left.red? and (sw == 1)
       
      # flip the color bits
      if !@left.nil?
        ll = @left.left
        if !ll.nil?
          if @left.red? and ll.red?
            copy(rot_right() || self)
            @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
      copy(rot_left() || self) if red? and !@right.nil? and @right.red? and (sw == 0)
       
      # flip the color bits
      r = @right
      if !@right.nil?
        rr = r.right
        if !rr.nil?
          if r.red? and rr.red?
            copy(rot_left() || self)
            @color = RedBlackTree::BLACK
            @left.color = RedBlackTree::RED unless @left.nil?
          end            
        end
      end
    end
  end
  
  # need a bunch of 'protected' accessors since 
  # in ruby one instance of an object can't see the instance variables of another instance
  attr_accessor :left, :right
  attr_writer :color
  
# ============================================================
# public methods
# ============================================================
public    
  def initialize(val = nil)
    @left      = nil
    @right     = nil
    @val       = val
    @color     = RedBlackTree::RED
  end
  
  def red?
    @color == RedBlackTree::RED
  end
  
  def black?
    @color == RedBlackTree::BLACK
  end
  
  def flip
    @color = (@color == RedBlackTree::RED) ? RedBlackTree::BLACK : RedBlackTree::RED
  end
  
  def to_s
    "[#{@val},#{@color}]"
  end

  # nice easy read-only accessors
  attr_reader :val, :color
  
  def find(key)
    if (key == @val) then
      self
    elsif key < @val
      @left.find(key) unless @left.nil?
    else
      @right.find(key) unless @right.nil?
    end
  end
       
  # inorder traversal using a lambda as the visitor
  def each_in_order(depth = 0, &block)
    @left.each_in_order(depth+1, &block) unless @left.nil?
    
    yield self, depth
    
    @right.each_in_order(depth+1, &block) unless @right.nil?
  end
  
  alias each each_in_order

  # main public insert
  def insert(*vals)
    vals.flatten.each { |val| rb_insert(val, 0) }
    @color = RedBlackTree::BLACK
  end
  
  alias << insert
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)
  root << nodelist
    
  # use a block directly 
  puts "Ruby B = #{root.map { |n,d| "(#{n.val}:#{n.color}:#{d})" }.join(', ')}"
  
  # do a find
  puts root.find(16)
end    

testRedBlackTree