#!/usr/bin/env ruby
# 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
class RedBlackTree
  RED   = 0
  BLACK = 1

  attr_reader :value, :color

  def initialize(value = nil)
    @left  = nil
    @right = nil
    @value = value
    @color = RED
  end

  def to_s
    "[#{@value}, #{@color == 0 ? 'RED' : 'BLACK'}]"
  end

  def find(key)
    result = nil
    if key == @value
      result = self
    elsif key < @value
      result = @left.find(key)  if @left
    else
      result = @right.find(key) if @right
    end
    result
  end

  def print(&block)
    @right.print_(1, &block) if @right
    yield self, 0
    @left.print_(1,  &block) if @left
  end

  def insert(value)
    insert_node(value, 0)
    @color = BLACK
  end

protected
  attr_accessor :left, :right
  attr_writer   :color

  def print_(depth, &block)
    @right.print_(depth + 1, &block) if @right
    yield self, depth
    @left.print_(depth + 1,  &block) if @left
  end

  def copy(node)
    @left  = node.left
    @right = node.right
    @value   = node.value
    @color = node.color
  end

  def insert_node_left(value, color)
    unless @left
      @left = RedBlackTree.new(value)
    else
      @left.insert_node(value, color)
    end
  end

  def insert_node_right(value,color)
    unless @right
      @right = RedBlackTree.new(value)
    else
      @right.insert_node(value, color)
    end
  end

  def left_rotate
    return nil unless @right

    # 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)
    node      = me.right
    me.right  = node.left
    node.left = me
    node
  end

  def right_rotate
    return nil unless @left

    # 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)
    node    = me.left
    me.left = node.right
    node.right = me
    node
  end

  def insert_node(value, color)
    # 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 && l.color == RED && r && r.color == RED
      @color  = RED
      l.color = r.color = BLACK
    end

    # go left or right depending on key relationship
    if value < @value then
      # recursively insert
      insert_node_left(value, 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 == RED && l && l.color == RED && color == 1
        node = right_rotate
        copy(node) if node
      end

      # flip the color bits
      l = @left
      if l && l.left && l.color == RED && l.left.color == RED
        node = right_rotate
        copy(node) if node
        @color = BLACK
        @right.color = RED if @right
      end
    else
      insert_node_right(value, 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 == RED && r && r.color == RED && color == 0
        node = left_rotate
        copy(node) if node
      end

      # flip the color bits
      r = @right
      if r && r.right && r.color == RED && r.right.color == RED
        node = left_rotate
        copy(node) if node
        @color = BLACK
        @left.color = RED if @left
      end
    end
  end
end


def testRedBlackTree
  #nodelist = [11,4,8,14,17,6,9,7,16,15,13,5,19,18,12,10,3,20,1]
  nodelist = [2, 3, 4, 5]
  root  = RedBlackTree.new(1)
  nodelist.each {|n| root.insert(n)}

  # update 2006-01-30 removed lambda example because it is equivalueent to the block example
  # which is more idiomatic Ruby

  # use a block directly
  root.print {|n,d| puts  "#{' '*(d*2)}#{n}" }
end

testRedBlackTree