Balanced (AVL) Binary Search Tree Implementation

Source.rb

#!/usr/bin/ruby

#require 'test/unit'
require 'test/unit/testcase'

class TestBST < Test::Unit::TestCase

  # ... tests omitted...

  def test_no_crashes
    b = BST.new
    (0..14).each do |n|
      now = Time.now
      b.randinit( 2**n )
      andThen = Time.now
      duration = andThen - now
      logb2 = Math.log( 2**n ) / Math.log( 2 )
      info = "n: #{2**n}, height: #{b.height}, logb2: #{logb2}, runtime: #{duration}s"
      # puts info
      #assert( ( b.height - logb2 ).abs <= 2.0, info )
    end
  end
end


class Node
  attr_reader :cached_height
  attr_accessor :key, :parent, :left, :right

  def initialize( key )
    @key = key
    @cached_height = 0
    @left = nil
    @right = nil
    @parent = nil
  end

  def find( key )
    return self if key == @key 
    return @left.find( key ) if @left and key < @key
    return @right.find( key ) if @right and key >= @key
    nil
  end

  def insert( key )
    if key >= @key
      if @right.nil?
        @right = Node.new( key )
        @right.parent = self
      else
        @right.insert( key )
      end

      update_height()
      rebalance()
      return @right
    else
      if @left.nil?
        @left = Node.new( key )
        @left.parent = self
      else
        @left.insert( key )
      end

      update_height()
      rebalance()
      return @left
    end
  end

  # returns nil
  def delete( key )
    if( key >= @key )
      @right.delete( key ) if @right
    else
      @left.delete( key ) if @left
    end
    update_height()
    rebalance()

    if @key == key
      if not @right.nil?
        node = @right.delete_min()
        @key = node.key
        update_height()
        rebalance()
      elsif not @left.nil?
        node = @left.delete_max()
        @key = node.key
        update_height()
        rebalance()
      elsif @parent
        @parent.left = nil if @parent.left.equal?( self )
        @parent.right = nil if @parent.right.equal?( self )
        @parent = nil
      end
    end
  end

  # Second try
  def to_s_recurse
    label = @key.to_s

    left, left_w = [], 1
    left, left_w = @left.to_s_recurse if @left
    right, right_w = [], 1
    right, right_w = @right.to_s_recurse if @right

    left << " " * left_w while left.length < right.length
    right << " " * right_w while right.length < left.length

    def mid( n ) (n+1)/2-1 end
    middle = 1
    w = left_w + middle + right_w
    g_mid = mid( w )
    header_begin_i = mid( left_w )
    header_end_i = left_w + middle + mid( right_w )
    label_begin_i = g_mid - mid( label.length )
    label_end_i = g_mid + label.length / 2

    branch_w = 1
    left_fix = (2 * ((header_begin_i + branch_w) - label_begin_i)) - ( 1 - w % 2 )
    right_fix = (2* (label_end_i - (header_end_i - branch_w)) - ( 1 - header_end_i % 2 ) )
    middle = [ middle, middle + left_fix, middle + right_fix ].max

    w = left_w + middle + right_w

    g_mid = mid( w )
    header_end_i = left_w + middle + mid( right_w )
    label_begin_i = g_mid - mid( label.length )
    label_end_i = g_mid + label.length / 2

    # "   ...345...   "
    above_branch = 1
    header = " " * (header_begin_i + above_branch)
    if label_begin_i > header_begin_i + above_branch
      header += "." * (label_begin_i - (header_begin_i + above_branch))
    end
    header += label
    if (header_end_i - above_branch) > label_end_i
      header += "." * (( header_end_i - above_branch ) - label_end_i )
    end
    header += " " * ( w - header_end_i )

    # "  /         \  "
    tree_branches = (" " * w)
    tree_branches[ header_begin_i ] = "/"
    tree_branches[ header_end_i ] = "\\"

    body = left.zip( right ).map { |l,r| l + " " * middle + r }

    lines = [ header, tree_branches ] + body
    return lines, w
  end

  def to_s
    lines, width = to_s_recurse
    return lines.join("\n")
  end

  def old_to_s
    l = if @left.nil? then "" else @left.to_s end
    r = if @right.nil? then "" else @right.to_s end

    @key.to_s + ": [ #{l}, #{r} ]"
  end

  def delete_min()
    if @left.nil?
      if @parent
        @parent.left = @right if @parent.left.equal?( self )
        @parent.right = @right if @parent.right.equal?( self )
      end

      @right.parent = @parent unless @right.nil?

      @right = nil
      @parent = nil

      return self
    end

    n = @left.delete_min()
    update_height()
    rebalance()
    n
  end

  def delete_max()
    if @right.nil?
      if @parent
        @parent.right = @left if @parent.right.equal?( self )
        @parent.left = @left if @parent.left.equal?( self )
      end

      @left.parent = @parent unless @left.nil?

      @left = nil
      @parent = nil

      return self
    end

    n = @right.delete_max()
    update_height()
    rebalance()
    n
  end

  # May not be self.
  def rotate_left
    c = @right

    @right = c.left
    c.left.parent = self if c.left

    c.left = self
    c.parent = @parent

    if @parent
      if @parent.right.equal?( self )
        @parent.right = c
      else
        @parent.left = c
      end
    end

    @parent = c

    update_height()
    c.update_height()
  end

  def rotate_right
    c = @left

    @left = c.right
    c.right.parent = self if c.right

    c.right = self
    c.parent = @parent

    if @parent
      if @parent.right.equal?( self )
        @parent.right = c
      else
        @parent.left = c
      end
    end

    @parent = c

    update_height()
    c.update_height()
  end

  def rebalance
    if height( @left ) >= height( @right ) + 2
      if( height( @left.left ) < height( @left.right ) )
        @left.rotate_left()
        rotate_right()
      else
        rotate_right()
      end
    elsif height( @right ) >= height( @left ) + 2
      if( height( @right.left ) > height( @right.right ) )
        @right.rotate_right()
        rotate_left()
      else
        rotate_left()
      end
    end
  end

  def update_height()
    @cached_height = [ height( @left ), height( @right ) ].max + 1
  end

private

  def height( node )
    return -1 if node.nil?
    node.cached_height
  end

end

class BST
  def initialize( vals = nil )
    @root = nil
    vals.each{ |v| self.insert( v ) } if not vals.nil?      
  end

  def randinit( n )
    n.to_i.times {
      self.insert( (rand() * 100 ).to_i )
    }
  end

  def height
    return -1 if not @root
    @root.cached_height
  end

  def insert( key )
    if @root.nil?
      @root = Node.new( key ) 
    else
      @root.insert( key )
    end

    # Root moved down in possible rebalance
    @root = @root.parent if @root.parent
  end

  def delete( key )
    @root.delete( key ) if @root

    @root = nil if @root.key == key

    # Root was moved down in a rebalance
    @root = @root.parent if @root and @root.parent
  end

  def find( key )
    return nil if not @root
    @root.find( key )
  end

  def to_s
    return "" if @root.nil?
    @root.to_s
  end
end

if ARGV.length == 1
  b = BST.new
  b.randinit( ARGV[0] )
  puts b
elsif ARGV.length > 1
  b = BST.new( ARGV.map { |i| i.to_i } )
  puts b
end