Binary Tree in Ruby
Today I thought I would have a play at implementing a binary tree in ruby and see if I could dig up any things I had not used before.
I first came up with a RSpec for it and set about building it and reading about the different methods you can use.
Here is the spec I came up with (This is the complete one, at the beginning they were all pending):
require "binary_tree"
describe BinaryTree do
before(:each) do
@tree = BinaryTree.new
@keys = %w{F B G A D I C E H}
@keys.each {|key| @tree.insert(key)}
end
it "should have a node for each key inserted" do
@keys.each do |key|
@tree.has?(key).should == true
end
end
it "should have the same number of keys as there are nodes" do
num = 0
@tree.traverse_inorder do |node|
num += 1
end
num.should eql(@keys.size)
end
it "should be tarversable in-order" do
expected = %w{A B C D E F G H I}
actual = []
@tree.traverse_inorder do |node|
actual << node.key
end
expected.should == actual
end
it "should be tarversable pre-order" do
expected = %w{F B A D C E G I H}
actual = []
@tree.traverse_preorder do |node|
actual << node.key
end
expected.should == actual
end
it "should be tarversable post-order" do
expected = %w{A C E D B H I G F}
actual = []
@tree.traverse_postorder do |node|
actual << node.key
end
expected.should == actual
end
it "should be able to return a node when given a key which it contains" do
@tree.insert("Z")
"Z".should == @tree.search("Z").key
end
it "should return nil if no matching key is found" do
nil.should == @tree.search("Z")
end
it "should be able to delete nodes"
it "should remain balanced following insertion and deletion of nodes"
end
As you can see there are still two to implement. The code that meets this specification is below.
# This class stores the tree and handles the insertion
# of new keys. Once the initial key has been inserted
# the insertion is handled by the node objects in the tree.
#
# Traversal of the tree is also controled by this class and
# can be pre-order, post-order or in-order by calling the
# relavant method.
#
# The BinaryTree should be able to store any object that includes
# the Comparable module.
#
# ====TODO
# - Impliment tree balancing
# - Impliment node deletion
#
class BinaryTree
# This class represents a single node in the tree.
# It holds a key and references to both its left
# and right sub trees.
#
# It is also able to determin if it is a leaf node
# using the .leaf? method.
#
class Node
attr_reader :key, :left, :right
# Create a Node and store the key as well
# as its left and right sub trees.
#
def initialize( key, left=nil, right=nil )
@key = key
@left, @right = left, right
end
# Insert a new key into the correct subtree based on the
# value of <i>self</i>.
#
# Duplicate keys are ignored.
#
def insert( key )
if key < @key
# insert left
@left.nil? ? @left = Node.new( key ) : @left.insert( key )
elsif key > @key
# insert right
@right.nil? ? @right = Node.new( key ) : @right.insert( key )
end # duplicate key, do nothing
end
# Test if <i>self</i> is a leaf node.
#
def leaf?
@left.nil? && @right.nil?
end
# Return the Key of <i>self</i> as a String
#
def to_s
@key.to_s
end
end
attr_reader :root
# Create a new BinaryTree
#
def initialize
@root = nil
end
# call-seq:
# insert(key) -> self
#
# Add a new key to <i>self</i>
#
# Eamples:
# tree = BinaryTree.new
# tree.insert(5).insert(2).insert(8).insert(1)
#
def insert( key )
if @root.nil?
@root = Node.new( key )
else
@root.insert( key )
end
self
end
# call-seq:
# traverse_inorder {|node| block } -> self
#
# Traverses <i>self</i> in-order.
#
# Calls <i>block</i> once for each node in <i>self</i>, passing that
# Node as a parameter.
#
# Eamples:
# tree = BinaryTree.new
# tree.insert(5).insert(2).insert(8).insert(1)
#
# tree.traverse_inorder { |node| print "#{node} " } # -> 1 2 5 8
#
def traverse_inorder( node=@root, &block )
return if node.nil?
traverse_inorder( node.left, &block )
yield node
traverse_inorder( node.right, &block )
end
# call-seq:
# traverse_preorder {|node| block } -> self
#
# Traverses <i>self</i> pre-order.
#
# Calls <i>block</i> once for each node in <i>self</i>, passing that
# Node as a parameter.
#
# Eamples:
# tree = BinaryTree.new
# tree.insert(5).insert(2).insert(8).insert(1)
#
# tree.traverse_preorder { |node| print "#{node} " } # -> 5 2 1 8
#
def traverse_preorder( node=@root, &block )
return if node.nil?
yield node
traverse_preorder( node.left, &block )
traverse_preorder( node.right, &block )
end
# call-seq:
# traverse_postorder {|node| block } -> self
#
# Traverses <i>self</i> post-order.
#
# Calls <i>block</i> once for each node in <i>self</i>, passing that
# Node as a parameter.
#
# Eamples:
# tree = BinaryTree.new
# tree.insert(5).insert(2).insert(8).insert(1)
#
# tree.traverse_postorder { |node| print "#{node} " } # -> 1 2 8 5
#
def traverse_postorder( node=@root, &block )
return if node.nil?
traverse_postorder( node.left, &block )
traverse_postorder( node.right, &block )
yield node
end
# call-seq:
# search(key) -> aNode
#
# Search <i>self</i> for <i>key</i> returning the Node if found
# or nil if not.
#
def search( key, node=@root )
return nil if node.nil?
if key < node.key
search(key, node.left)
elsif key > node.key:
search(key, node.right)
else
return node
end
end
# Test <i>self</i> for <i>key</i>.
#
def has?( key )
not self.search(key).nil?
end
end
Download all the code from here: binary_tree.tgz and have a play. All the comments have been written with RDoc in mind so the documentation that is generated is quite nice. This example may give someone a helpful working example of RSpec and RDoc.