# $Id: rbtree.rb,v 1.3 2001/10/22 20:41:08 ser Exp $
# Author: Sean Russell <ser@germane-software.com>
#         http://www.germane-software.com/~ser

## RedBlackTree
# An ordered tree object.  Items put into this tree must be Comparable (<=>).
# This class mixes in Enumerable, and optimized versions of min and
# max are supplied.  All other Enumerable methods are unoptimized, and
# performance for these may be worse than an Array.  min() and max() are only
# optimized for the default case (IE, no supplied block).  Note that these
# limitations are due to the fact that this is an *ordered* tree, and any
# attempt to perform operations on the tree that can not take advantage of
# this order will suffer.
class RedBlackTree
	include Enumerable

	RED = "R"
	BLACK = "B"

	attr_reader :root

	def initialize
		@root = nil
	end

	# A synonym for add()
	def << object
		add object
	end

	# Adds a Comparable object to this tree
	def add object
		x = insert object # insert the object, sorted, into the tree
		x.color = RED
		rot = proc { |x, y, left, right|
			y = x.parent.parent.send(right)
			if !y.nil? and y.color == RED
				x.parent.color = BLACK
				y.color = BLACK
				x.parent.parent.color = RED
				x = x.parent.parent
			else
				if x == x.parent.send(right)
					x = x.parent
					self.send("#{left}_rotate", x)
				end
				x.parent.color = BLACK
				x.parent.parent.color = RED
				self.send("#{right}_rotate", x.parent.parent)
			end
		}

		while x != @root and x.parent.color == RED
			if x.parent == x.parent.parent.left
				y = x.parent.parent.right
				rot.call(x,y,"left", "right")
			else
				y = x.parent.parent.left
				rot.call(x,y,"right", "left")
			end
		end
		@root.color = BLACK
	end

	# yields for each item in the tree
	def each
		traverse {|item, count|
			yield item.value unless item.nil?
		}
	end

	# Finds an object in the tree.  If a block is provided, it will be yielded
	# to with each node that is touched on the way to the item.
	def search value
		curr = @root
		while true
			return nil if curr.nil?
			yield curr.value if block_given?
			case curr.value <=> value
			when 0
				return curr.value
			when -1
				curr = curr.left
			else
				curr = curr.right
			end
		end
		nil
	end

	# Finds a path to a given value, as an array of nodes leading to the item.
	def path_to value
	  rv = []
	  search(value) { |val| rv << val }
	  rv
	end
	
	def min(&block)
		if block_given?
			return super.min(&block)
		end
		curr = @root
		while true
			return curr if curr.left.nil?
			curr = curr.left
		end
	end

	def max
		if block_given?
			return super.max(&block)
		end
		curr = @root
		while true
			return curr if curr.right.nil?
			curr = curr.right
		end
	end

	def traverse object=@root
		depth_first = proc { |node, count|
			yield(node, count)
			depth_first.call(node.left, count+1) unless node.left.nil?
			depth_first.call(node.right, count+1) unless node.right.nil?
		}
		depth_first.call object, 0
	end

	def max_depth_of child
		max_depth = 0
		traverse(child) { |node, count| max_depth = count if count > max_depth }
		max_depth
	end

	private
	def left_rotate x
		y = x.right
		x.right = y.left
		y.left.parent = x if y.left != nil
		y.parent = x.parent
		if x.parent == nil then @root = y
		else
			if x == x.parent.left then x.parent.left = y
			else x.parent.right = y
			end
		end
		y.left = x
		x.parent = y
	end

	def right_rotate x
		y = x.left
		x.left = y.right
		y.right.parent = x if y.right != nil
		y.parent = x.parent
		if x.parent == nil then @root = y
		else
			if x == x.parent.right then x.parent.right = y
			else x.parent.left = y
			end
		end
		y.right = x
		x.parent = y
	end

	def insert object
		object = Node.new( object )
		if @root == nil
			@root = object
			return object
		end

		insert_helper @root, object
		object
	end

	def insert_helper parent, child
	  if child.value < parent.value
		  if parent.left == nil
			  parent.left = child
			  child.parent = parent
		  else
			  insert_helper( parent.left, child )
		  end
	  else
		  if parent.right == nil
			  parent.right = child
			  child.parent = parent
		  else
			  insert_helper( parent.right, child )
		  end
	  end
	end

	class Node
		attr_accessor :left, :right, :color, :parent
		attr_reader :value
		def initialize value
			@value = value
			@color = RED
			@left = @right = @parent = nil
		end
	end
end
