Binary Search Tree in Ruby

binary_search_tree.rb

=begin

3,5,2,7,1

  3
 2 5
1   7

=end

class Person
  attr_accessor :uid, :first_name, :last_name, :age

  def initialize(uid, first_name, last_name, age)
    self.uid, self.first_name, self.last_name, self.age = uid, first_name, last_name, age
  end
end

@people = []
@people << Person.new(3,'anthony','stark',44)
@people << Person.new(5,'peter','parker',18)
@people << Person.new(2,'mathew','murdock',28)
@people << Person.new(7,'bruce','banner',32)
@people << Person.new(4,'stephen','strange',52)


class BinarySearchTree
  attr_accessor :left, :right, :value

  DEFAULT_COMPARISON = lambda{|a,b| a <=> b }

  def initialize(&custom_comparison)
    @comparison = block_given? ? custom_comparison : DEFAULT_COMPARISON
  end

  def to_s
    str = " #{value.inspect} " if value
    str << "=> [" if (left || right)
    str << "#{left}" if left
    str << "," if left && right
    str << "#{right}" if right
    str << "]" if (left || right)
    return str
  end

  def depth
    depth_left = left.nil? ? 0 : left.depth
    depth_right = right.nil? ? 0 : right.depth
    depth_children = depth_left >= depth_right ? depth_left : depth_right
    return depth_children + 1
  end

  def add(new_value)
    if new_value
      if self.value.nil?
        puts "New"
        self.value = new_value
      else
        comparison = @comparison.call(new_value, value)        
        if comparison == 0
          puts "Replace"
          self.value = new_value
        elsif comparison > 0
          puts "Right"
          self.right ||= BinarySearchTree.new(&@comparison)
          self.right.add(new_value)
        elsif comparison < 0
          puts "Left"
          self.left ||= BinarySearchTree.new(&@comparison)
          self.left.add(new_value)
        else
          raise "invalid state"
        end
      end
    end
    return self
  end

  def <<(new_value)
    add(new_value)
  end

  def insert(*values)
    values.each{|new_value| add(new_value) }
    return self
  end

  def search(search_value, subtree_result=false)
    if search_value.nil? || value.nil?
      puts "Nothing"
    else
      comparison = @comparison.call(search_value, value)
      if comparison == 0
        puts "Found"
        return subtree_result ? self : self.value
      elsif comparison > 0
        puts "Right"
        return right.search(search_value, subtree_result) if right
      elsif comparison < 0
        puts "Left"
        return left.search(search_value, subtree_result) if left
      else
        raise "invalid state"
      end
    end
    return nil
  end

  def traverse(&callback)
    left.traverse(&callback) if left
    callback.call(value)
    right.traverse(&callback) if right    
  end

end

class TreeMap
  class Entry
    attr_accessor :key, :value

    def initialize(key, value) 
      self.key, self.value = key, value
    end
  end

  attr_accessor :left, :right, :entry

  DEFAULT_COMPARISON = lambda{|a,b| a <=> b }

  def initialize(&custom_comparison)
    @comparison = block_given? ? custom_comparison : DEFAULT_COMPARISON
  end

  def to_s
    str = " #{entry.value.inspect} " if entry
    str << "=> [" if (left || right)
    str << "#{left}" if left
    str << "," if left && right
    str << "#{right}" if right
    str << "]" if (left || right)
    return str
  end

  def depth
    depth_left = left.nil? ? 0 : left.depth
    depth_right = right.nil? ? 0 : right.depth
    depth_children = depth_left >= depth_right ? depth_left : depth_right
    return depth_children + 1
  end

  def put(key, value)
    raise ArgumentError.new("key can not be nil") if key.nil?

    if entry.nil?
      puts "New"
      self.entry = Entry.new(key,value)
    else
      comparison = @comparison.call(key, entry.key)        
      if comparison == 0
        puts "Replace"
        self.entry.value = value
      elsif comparison > 0
        puts "Right"
        self.right ||= TreeMap.new(&@comparison)
        self.right.put(key, value)
      elsif comparison < 0
        puts "Left"
        self.left ||= TreeMap.new(&@comparison)
        self.left.put(key, value)
      else
        raise "invalid state"
      end
    end

    return self
  end

  def get(key)
    if entry.nil? || key.nil?
      puts "Nothing"
    else
      comparison = @comparison.call(key, entry.key)
      if comparison == 0
        puts "Found"
        return entry.value
      elsif comparison > 0
        puts "Right"
        return right.get(key) if right
      elsif comparison < 0
        puts "Left"
        return left.get(key) if left
      else
        raise "invalid state"
      end
    end
    return nil
  end

  def [](key)
    get(key)
  end

  def []=(key, value)
    put(key, value)
  end

  def traverse(&callback)
    left.traverse(&callback) if left
    callback.call(entry.key, entry.value)
    right.traverse(&callback) if right    
  end

  def to_h
    hash = {}
    traverse{|key,value| hash[key] = value}
    return hash
  end

end