Right balanced AVL tree for minimizing binary comparisons in order to find top k items.

gistfile1.rb

$counter = 0
 
class RightBalancedTree
  attr_reader :root
 
  def initialize(top_x=25, node_type=AVLTreeNode)
    @root = nil
    @node_type = node_type
    @top_x
  end
 
  def add(value)
    if @root
      side = @root.add_child(@node_type.new(value, nil, nil))
      if side == :left
        @root.left = nil
      else
        if @root.right.count > @top_x
          node = @root.right.kill_leftmost_child
          if node
            @root.value = node.value
          else
            @root = @root.right
          end
        end
        @root.right.balance
      end
    else
      @root = @node_type.new(value, nil, nil)
    end
    
    nil
  end
end
 
 
class NilClass
  def height
    -1
  end
 
  def count
    0
  end
  
  def right
    nil
  end
  
  def left
    nil
  end
  
  def balance_factor
    0
  end
end
 
 
class AVLTreeNode
  attr_accessor :value, :left, :right
 
 
  def initialize(value, left, right)
    @value = value
    @left = left
    @right = right
  end
 
 
  def add_child(node)
    if self > node
      if @left
        @left.add_child(node)
      else
        @left = node
        :left
      end
    else
      if @right
        @right.add_child(node)
      else
        @right = node
        :right
      end
    end
  end
 
 
  def height
    1 + [@left.height, @right.height].max
  end
 
 
  def count
    1 + @left.count + @right.count
  end
 
 
  def balance_factor
    @left.height - @right.height
  end
 
 
  def minimum_child
    if @left
      @left.minimum_child
    else
      self
    end
  end
 
 
  def <=>(node)
    $counter += 1
    @value <=> node.value
  end
 
  def ==(node)
    $counter += 1
    @value == node.value
  end
 
  def <(node)
    $counter += 1
    @value < node.value
  end
 
  def >(node)
    $counter += 1
    @value > node.value
  end
 
  def >=(node)
    $counter += 1
    @value >= node.value
  end
 
  def <=(node)
    $counter += 1
    @value <= node.value
  end
 
  def to_s
    %{TreeNode(#{@value}, #{@left.to_s}, #{@right.to_s})}
  end
 
 
  def rotate_LL
    tmp = @right
    @right = @left
    @left = @right.left
    @right.left = @right.right
    @right.right = tmp
 
    tmp = @value
    @value = @right.value
    @right.value = tmp
  end
 
 
  def rotate_LR
    @left.rotate_RR
    rotate_LL
  end
  
  
  def rotate_RL
    @right.rotate_LL
    rotate_RR
  end
 
 
  def rotate_RR
    tmp = @left
    @left = @right
    @right = @left.right
    @left.right = @left.left
    @left.left = tmp
 
    tmp = @value
    @value = @left.value
    @left.value = tmp
  end
 
 
  def balance
    if balance_factor > 1
      if @left.balance_factor > 0
        rotate_LL
      else
        rotate_LR
      end
    elsif balance_factor < -1
      if @right.balance_factor < 0
        rotate_RR
      else
        rotate_RL
      end
    end
  end
  
  
  def prune
    if @left
      @left.prune
      @left = nil
    end
    if @right
      @right.prune
      @right = nil
    end
  end
  
  
  def kill_leftmost_child
    if @left
      if @left.left
        @left.kill_leftmost_child
      else
        tmp = @left
        @left = nil
        tmp
      end
    else
      nil
    end
  end
 
end