On 12/18/07, Adam Shelly wrote:
>I added test_remove_multiple_nodes - there are definitely a few
>mistakes in the current remove :)

Seems the only mistake was that remove was expecting AVLTree#each to return the node traversed, but it was instead returning the node's data. I modified each to accept the name of a function to be passed to itself, which defaults to data. I made an "identity" method called node, and the passed it in when each is called within remove.

For my test, each and to_a must now accept symbols denoting the type of traversal to be performed, including inorder (the current traversal), preorder, postorder, and level-by-level.

==>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(*args,&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(returns=:data,&iter)
      @left.each(&iter)
      iter[self.send(returns)]
      @right.each(&iter)
    end
    
    def node
      self
    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(:node) 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(:node) 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(*args,&iter)
    @root.each(*args,&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<==


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_alternate_traversals
    items = [3,2,4,1,5]
    items.each {|el| @tree << el}
    
    preorder_result = [3,2,1,4,5]
    assert_equal(@tree.to_a(:preorder),preorder_result)
    @tree.each(:preorder) {|el| assert_equal(preorder_result.shift,el)}
    
    inorder_result = [1,2,3,4,5]
    assert_equal(@tree.to_a(:inorder),inorder_result)
    @tree.each(:inorder) {|el| assert_equal(inorder_result.shift,el)}
    
    postorder_result = [1,2,5,4,3]
    assert_equal(@tree.to_a(:postorder),postorder_result)
    @tree.each(:postorder) {|el| assert_equal(postorder_result.shift,el)}
    
    bylevel_result = [3,2,4,1,5]
    assert_equal(@tree.to_a(:by_level),bylevel_result)
    @tree.each(:by_level) {|el| assert_equal(bylevel_result.shift,el)}
  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
_END_



      ____________________________________________________________________________________
Never miss a thing.  Make Yahoo your home page. 
http://www.yahoo.com/r/hs