Binary Search Tree sauce Ruby (part 1)

This post is the first of two focusing on the implementation of a binary search tree in Ruby. This first part will deal with adding elements to a binary search tree. The second part will cover binary tree traversals.

Let’s jump straight in! After writing the unit tests, I have a first go at the code (I chose to go for a linked implementation rather than an array based one). I end up with this:

class BinaryTree
  def add(value)
    if @root.nil?
      @root = BinaryTreeNode.new(value)
    else
      @root.add(value)
    end
  end
end

class BinaryTreeNode
  def initialize(value)
    @value = value
  end

  def add(value)
    if value < @value
      if @left.nil?
        @left = BinaryTreeNode.new(value)
      else
        @left.add(value)
      end
    else
      if @right.nil?
        @right = BinaryTreeNode.new(value)
      else
        @right.add(value)
      end
    end
  end
end

Green light, the unit tests pass. One thing is annoying me though: code duplication! In three places, once in BinaryTree and twice in BinaryTreeNode, I have:

if field.nil?
  field = BinaryTreeNode.new(value)
else
  field.add(value)
end

Wouldn’t it be great to have a method that writes this piece of code for a particular field? Metaprogramming, here we come! I will put this method in a module so that it can be included in both classes.

module BinaryTreeHelper
  private
  def add_or_create_node(field, value)
    if instance_variable_get(field).nil?
      instance_variable_set(field,
                            BinaryTreeNode.new(value))
    else
      instance_variable_get(field).add(value)
    end
  end
end

Now I can refactor my classes:

class BinaryTree
  include BinaryTreeHelper

  def add(value)
    add_or_create_node(:@root, value)
  end
end

class BinaryTreeNode
  include BinaryTreeHelper

  def initialize(value)
    @value = value
  end

  def add(value)
    add_or_create_node(value < @value ? :@left : :@right,
                       value)
  end
end

Pretty nice, no?