From: Michal Suchanek <hramrach@...>
Date: 2007-12-15T08:42:02+09:00
Subject: Re: [QUIZ] Programmer Ping-Pong (#150)

This one is really a tree for once.

#!/usr/bin/env ruby -wKU

class TreeNode

  attr_accessor :left, :right, :data, :parent

  def initialize obj
    @left=nil
    @right=nil
    @data=obj
    @parent=nil
  end

  def attach_left node
    @left = node
    node.parent = self
  end

  def attach_right node
    @right = node
    node.parent = self
  end

  def height
    max( (left.height rescue 0) , (right.height rescue 0) )+1
  end

  def max *args
    args.inject{|m,n| n>m ? n : m}
  end

  def << node
    left ? (right ? (left << node) : (attach_right node)) : (attach_left node)
  end

  def include? obj
    (@data == obj) ||
      (@left.include? obj rescue false) ||
      (@right.include? obj rescue false)
  end

  def length
    len = 1
    len += @left.length if left
    len += @right.length if right
  end

end

class AVLTree

 def initialize
  @root = nil
 end

 def empty?
  ! @root
 end

 def include?(obj)
  (! empty?) && (@root.include? obj)
 end

 def <<(obj)
  if empty?
    @root = TreeNode.new obj
  else
    @root << TreeNode.new(obj)
  end
  self
 end

 def height
  empty? ? 0 : @root.height
 end

end

if $0 == __FILE__
  require "test/unit"

  class TestAVLTree < Test::Unit::TestCase
    def setup
      @tree = AVLTree.new
    end


    def test_tree_membership
      assert_equal(true,  @tree.empty?)
      assert_equal(false, @tree.include?(3))

      @tree << 3

      assert_equal(false, @tree.empty?)
      assert_equal(true,  @tree.include?(3))
    end

    def test_tree_height_of_one_or_two_nodes_is_N
      @tree << 5
      assert_equal 1, @tree.height
      @tree << 6
      assert_equal 2, @tree.height     #changed from 1
    end


    def test_tree_insertion
      assert_equal(true, @tree.empty?)
      assert_equal(false, @tree.include?(3))
      assert_equal(false, @tree.include?(5))

      @tree << 3
      @tree << 5

      assert_equal(false, @tree.empty?)
      assert_equal(true, @tree.include?(5))
      assert_equal(true, @tree.include?(3))
    end

    def test_tree_include_many
      0.upto(10) do |i|
        assert(! @tree.include?(i), "Tree should not include #{i} yet.")
        @tree << i
        0.upto(i) do |j|
          assert( @tree.include?(j), "Tree should include #{j} already.")
        end
      end
    end

    def test_tree_traverse
      ary = [3,5,17,30,42,54,1,2]
      ary.each{|n| @tree << n}
      traversal = []
      @tree.each{|n| traversal << n}
      ary.each{|n| assert traversal.include?(n), "#{n} was not visited
in tree."}
    end

  end
end