1class AVLTree
  2  attr_reader :root
  3
  4  def push(item)
  5    if @root
  6      @root = @root.push(item)
  7    else
  8      @root = Node.new(item)
  9    end
 10  end
 11
 12  def size
 13    return 0 unless @root
 14    @root.size
 15  end
 16
 17  def height
 18    return 0 unless @root
 19    @root.height
 20  end
 21
 22  alias_method :add, :push
 23
 24  class Node
 25    attr_accessor :data, :left, :right
 26
 27    def initialize(data)
 28      @data = data
 29    end
 30
 31    def push(item)
 32      if (item <=> @data) == -1
 33        @left.push(item) if @left
 34        @left = Node.new(item) unless @left
 35      else
 36        @right.push(item) if @right
 37        @right = Node.new(item) unless @right
 38      end
 39      re_balance
 40    end
 41
 42    def size
 43      result = 1
 44      result += @left.size if @left
 45      result += @right.size if @right
 46      result
 47    end
 48
 49    def height
 50      result = 1
 51      result += left_height >= right_height ? left_height : right_height
 52      result
 53    end
 54
 55    def re_balance
 56      if balance_factor < -1
 57        if @right.balance_factor == 1
 58          right_left_case
 59        else
 60          right_right_case
 61        end
 62      elsif balance_factor > 1
 63        if @left.balance_factor == -1
 64          left_right_case
 65        else
 66          left_left_case
 67        end
 68      else
 69        self
 70      end
 71    end
 72
 73    def balance_factor
 74      (left_height - right_height) || 0
 75    end
 76
 77    private
 78
 79    def right_left_case
 80      right = @right
 81      @right = right.left
 82      @right.right = right
 83      right.left = nil
 84      right_right_case
 85    end
 86
 87    def right_right_case
 88      @right.left = self
 89      new_root = @right
 90      @right = nil
 91      new_root
 92    end
 93
 94    def left_right_case
 95      left = @left
 96      @left = @left.right
 97      @left.left = left
 98      left.right = nil
 99      left_left_case
100    end
101
102    def left_left_case
103      @left.right = self
104      new_root = @left
105      @left = nil
106      new_root
107    end
108
109    def left_height
110      @left ? @left.height : 0
111    end
112
113    def right_height
114      @right ? @right.height : 0
115    end
116  end
117end