Re: [QUIZ] Programmer Ping-Pong (#150)

On 12/18/07, James Gray <[e-mail]> wrote:
> I took a first stab at the removal code.  I implemented more than was
> strictly needed to make the test pass, because I haven't helped out
> much yet.  However, I may have made mistakes here and we should
> definitely add more tests to double-check me.

I added test_remove_multiple_nodes - there are definitely a few
mistakes in the current remove :)
>
> Before that though, I wanted to see if I could distract us with a
> little interface test.
>

Nice test - here's a fairly simple solution.
I also implemented some caching of @height results, in the hopes of
getting closer to that O(log N)

-Adam
==> avl_tree.rb <==
class AVLTree

   include Enumerable

   # Need something smarter than nil for external nodes
   class ExternalNode
     attr_accessor :parent
     def initialize(parent)
       @parent = parent
     end
     def include?(_)     false              end
     def <<(_)           raise RuntimeError end
     def height;         0                 end
     def balance_factor; 0                 end
     def left;           self               end
     def right;          self               end
     def each(&iter)                        end
     def left=(node)
       raise(NotImplementedError,
             "external node of #{@parent}")
     end
     def right=(node)
       raise(NotImplementedError,
             "external node of #{@parent}")
     end
     def to_s; '' end
     def to_a; [] end
   end

   class Node
     attr_accessor :data, :parent
     attr_reader :left, :right

     def initialize obj, sortblock
       @parent = nil
       @data   = obj
       @left   = ExternalNode.new(self)
       @right  = ExternalNode.new(self)
       @height = 1
       @compare = sortblock
     end

     def left=(node)
       @height = nil
       @left = node
       node.parent = self
     end

     def right=(node)
       @height = nil
       @right = node
       node.parent = self
     end

     def each(&iter)
       @left.each(&iter)
       iter[data]
       @right.each(&iter)
     end

     def height
       @height || ( [ @left.height, @right.height ].max + 1)
     end

     def << node
       case @compare[node.data,@data]
       when -1
        if Node === @left
           @left << node
         else
          self.left = node
         end
       when 0
         return self             # no dups
       when +1
         if Node === @right
           @right << node
         else
           self.right = node
         end
       end
       rebalance if balance_factor.abs > 1
       @height = nil
     end

     def remove obj
       @height = nil
       case @compare[obj,@data]
       when -1
         if Node === @left
           @left.remove(obj)
         else
           nil
         end
       when 0
         if Node === @left
           largest = nil
           AVLTree.new(@left).each do |n|
             largest = n if largest.nil? or n.data > largest.data
             largest = n if largest.nil? or n.data > largest.data
           end
           self.data = largest.data
           self.left = largest.left
         elsif Node === @right
           smallest = nil
           AVLTree.new(@right).each do |n|
             smallest = n if smallest.nil? or n.data < smallest.data
             smallest = n if smallest.nil? or n.data < smallest.data
           end
           self.data  = smallest.data
           self.right = smallest.right
         else
           parent.send :replace_child, self, ExternalNode.new(parent)
         end
         path = self
         while Node === (path = path.parent)
           path.rebalance if path.balance_factor.abs > 1
         end
         self
       when +1
         if Node === @right
           @right.remove(obj)
         else
           nil
         end
       end
     end

     def include? obj
       case obj <=> @data
       when -1 : @left.include?(obj)
       when  0 : true
       when +1 : @right.include?(obj)
       end
     end

     def to_a
       result = @left.to_a
       result << @data
       result.concat @right.to_a
       result
     end

     def [](idx)
        if idx < (leftheight = @left.height)
          @left[idx]
        elsif (idx== leftheight)
          @data
        elsif (idx-=(leftheight+1)) < @right.height
          @right[idx]
        end
     end

     def to_s
       bf = case balance_factor <=> 0
            when -1 : '-' * -balance_factor
            when  0 : '.'
            when  1 : '+' * balance_factor
            end
       "[#{left} "+
         "(#{@data}{#{height}#{bf}}^#{parent.data})"+
         " #{right}]"
     end

     protected

     def balance_factor
       @right.height - @left.height
     end

     def rotate_left
       my_parent, from, to = @parent, self, @right
       temp = @right.left
       @right.left = self
       self.right = temp
       my_parent.send :replace_child, from, to
       to.parent = my_parent
     end

     def rotate_right
       my_parent, from, to = @parent, self, @left
       temp = @left.right
       @left.right = self
       self.left = temp
       my_parent.send :replace_child, from, to
       to.parent = my_parent
     end

     def rebalance
       if (bf = balance_factor) > 1 # right is too high
         if @right.balance_factor < 0
           # double rotate right-left
           # - first the right subtree
           @right.rotate_right
         end
         rotate_left             # single rotate left
       elsif bf < -1             # left must be too high
         if @left.balance_factor > 0
           # double rotate left-right
           # - first force left subtree
           @left.rotate_left
         end
         rotate_right            # single rotate right
       end
     end

     def replace_child(from, to)
       if from.eql? @left
         @left = to
       elsif from.eql? @right
         @right = to
       else
         raise(ArgumentError,
               "#{from} is not a branch of #{self}")
       end
     end

   end

   def initialize(root = nil, &block)
     @root = root
     if block
      raise(ArgumentError,
         "Block argument for #{self.class.name} must" +
         " take 2 arguments and act as sort function"
         ) unless block.arity == 2
     else
       block = proc{|a,b| a<=>b}
     end
     @compare = block
   end

   def empty?
     @root.nil?
   end

   def include?(obj)
     empty? ? false : @root.include?(obj)
   end

   def <<(obj)
     raise(ArgumentError,
           "Objects added to #{self.class.name} must" +
           " respond to <=>"
           ) unless obj.respond_to?(:<=>)

     if empty?
       @root = Node.new(obj, @compare)
       @root.parent = self
     else
       @root << Node.new(obj, @compare)
     end
     self
   end

   def remove(obj)
     @root.remove(obj) unless empty?
   end

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

   def [](idx)
     @root[idx]
   end

   def to_a
     empty? ? [] : @root.to_a
   end

   def each(&iter)
     @root.each(&iter) unless empty?
   end

   def to_s
     empty? ? "[]" : @root.to_s
   end

   # Indicates that parent is root in to_s
   def data; '*'; end

   protected

   def replace_child(from, to)
     if @root.eql? from
       @root = to
     else
       raise(ArgumentError,
             "#{from} is not a branch of #{self}")
     end
   end

end
__END__

==> test_avl_tree.rb <==
#!/usr/bin/env ruby -wKU

require "test/unit"
require "avl_tree"

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

   ##################################################
   # Membership tests
   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_should_allow_more_than_one_element
     @tree << 3
     @tree << 4

     assert(@tree.include?(4), "4 not in #{@tree}")
     assert(@tree.include?(3), "3 not in #{@tree}")
   end

   def test_tree_include_many
     0.upto(10) do |i|
       assert_equal(false, @tree.include?(i),
                    "Tree should not include #{i} yet.")
       @tree << i
       0.upto(i) do |j|
         assert_equal(true, @tree.include?(j),
                      "Tree should have 0..#{i},"+
                      " where's #{j}? ")
       end
     end
   end

   # This sits at the intersection of membership
   # and height tests.  We know one node has height 1,
   # and two nodes has height 2, so if we insert one
   # object twice and the height is 1, there must
   # only be one node in the tree.
   def test_tree_does_not_keep_duplicates
     @tree << 'a'
     @tree << 'a'
     assert_equal 1, @tree.height, "one node: #{@tree}"
   end

   ##################################################
   # Height tests
   def test_tree_height_of_one_or_two_nodes_is_N
     @tree << 5
     assert_equal 1, @tree.height, "one node: #{@tree}"
     @tree << 6
     assert_equal 2, @tree.height, "two nodes: #{@tree}"
   end

   def test_tree_height_of_three_nodes_is_two
     @tree << 5
     @tree << 6
     @tree << 7
     assert_equal 2, @tree.height, @tree.to_s
   end

   # RobB: The more precise limit given in [Knuth] is
   # used rather than that from [Wikipedia]
   def test_tree_growth_limit_is_1pt44_log_N
     (1..10).each do |i|
       @tree << i
       limit = ((1.4405 *
                 Math::log(i+2.0)/Math::log(2.0)
                 ) - 0.3277).ceil
       assert(@tree.height <= limit,
              "Tree of #{i} nodes is too tall by" +
              " #{@tree.height - limit}" +
              " #{@tree}")
     end
   end

   def test_balances_left
     4.downto(1) { |i| @tree << i }
     assert(@tree.height < 4,
            "expected tree height #{@tree.height} < 4")
   end

   def test_balances_right
     1.upto(4) { |i| @tree << i }
     assert(@tree.height < 4,
            "expected tree height #{@tree.height} < 4")
   end

   def test_non_sequential_insertion__part_1
     items = [ 1, 3, 2 ]
     items.each do |i|
       @tree << i
     end
     items.each do |i|
       assert_equal(true, @tree.include?(i),
                    "where is #{i}? ")
     end
   end

   def test_non_sequential_insertion__part_2
     items = [ 3, 1, 2 ]
     items.each do |i|
       @tree << i
     end
     items.each do |i|
       assert_equal(true, @tree.include?(i),
                    "where is #{i}? ")
     end
   end

   ##################################################
   # Access tests (getting data back out)

   # RobB: this tests too much at one time; I sorted ary.
   def test_tree_traverse
     ary = [ 3, 5, 17, 30, 42, 54, 1, 2 ].sort

     ary.each { |n| @tree << n }
     traversal = []
     @tree.each { |n| traversal << n }

     assert_equal(ary.size, traversal.size)

     ary.each do |n|
       assert_equal(true, traversal.include?(n),
                    "#{n} was not visited in tree.")
     end
   end

   def test_tree_find
     [1,2,3,4].each{|n| @tree<<n}
     assert_equal(1, @tree.find{|v|v>0} )
     assert_equal(2, @tree.find{|v|v%2==0} )
   end

   def test_sequential_access
     items = [ 50, 17, 72 ]
     items.each { |n| @tree << n }
     items.sort.each_with_index do |e,i|
       assert_equal(e, @tree[i],
                    "@tree[#{i}] should be like " +
                    "#{items.inspect}.sort[#{i}]")
     end
   end

   # [Knuth] p.473: "The problem of deletion can be solved
   # in O(log N) steps if we approach it correctly."
   def test_remove_node
     @tree << 314
     @tree.remove(314)
     assert_equal(false, @tree.include?(314),
                  '314 still in the tree')
     end

    def test_remove_multiple_nodes
     items = [ 50, 17, 72, 45, 43, 23 ]
     items.each { |n| @tree << n }
     puts @tree, @tree.height
     @tree.remove(50)
     assert_equal(false, @tree.include?(50),
                  '50 still in the tree')
     @tree.remove(72)
     assert_equal(false, @tree.include?(72),
                  '72 still in the tree')
     @tree.remove(45)
     assert_equal(false, @tree.include?(45),
                  '45 still in the tree')
     assert_equal(2, @tree.height)  #tree should have 3 items, height = 2
    end



   ##################################################
   # Interface tests

   def test_custom_comparison_code
     rev_tree = AVLTree.new { |a, b| b <=> a }
     values   = [3, 2, 1]
     values.each { |v| rev_tree << v }
     rev_tree.each { |v| assert_equal(values.shift, v) }

     len_tree = AVLTree.new { |a, b| a.length <=> b.length }
     values   = %w[3 22 111]
     values.each { |v| len_tree << v }
     len_tree.each { |v| assert_equal(values.shift, v) }
   end

end
=begin
[Knuth] Knuth, Donald Ervin,
     The Art of Computer Programming, 2nd ed.
     Volume 3, Sorting and Searching,
     section 6.2.3 "Balanced Trees", pp. 458-478
[Wikipedia]
     AVL Tree, http://en.wikipedia.org/wiki/AVL_tree
=end
__END__