Commit 6b2cc28

@@ -1,10 +1,10 @@
 class AVLTree
-  def initialize
-  end
+  attr_reader :root
 
   def push(item)
+    p "insert #{item}"
     if @root
-      @root.push(item)
+      @root = @root.push(item)
     else
       @root = Node.new(item)
     end
@@ -23,6 +23,8 @@ class AVLTree
   alias_method :add, :push
 
   class Node
+    attr_accessor :data, :left, :right
+
     def initialize(data)
       @data = data
     end
@@ -35,6 +37,7 @@ class AVLTree
         @right.push(item) if @right
         @right = Node.new(item) unless @right
       end
+      re_balance
     end
 
     def size
@@ -46,10 +49,36 @@ class AVLTree
 
     def height
       result = 1
-      left_height = @left ? @left.height : 0
-      right_height = @right ? @right.height : 0
       result += left_height >= right_height ? left_height : right_height
       result
     end
+
+    def re_balance
+      p balance_factor
+      if balance_factor < -1
+        p "right right case"
+        @right.left = self
+        new_root = @right
+        @right = nil
+        new_root
+      else
+        self
+      end
+    end
+
+    private
+
+    def balance_factor
+      p "left: #{left_height}, right: #{right_height}"
+      (left_height - right_height) || 0
+    end
+
+    def left_height
+      @left ? @left.height : 0
+    end
+
+    def right_height
+      @right ? @right.height : 0
+    end
   end
 end

@@ -57,9 +57,29 @@ describe AVLTree do
     end
   end
 
-  context "when adding an item that would cause the tree to be unbalanced" do
-    it "should re-balance it self" do
-      
+  context "when the tree is unbalanced" do
+    context "with a right-right case" do
+      before :each do
+        sut.add("a")
+        sut.add("b")
+        sut.add("c")
+      end
+
+      it "should re-balance it self" do
+        sut.height.should == 2
+      end
+
+      it "should have a new root" do
+        sut.root.data.should == "b"
+      end
+
+      it "should change the left side" do
+        sut.root.left.data.should == "a"
+      end
+
+      it "should change the right side" do
+        sut.root.right.data.should == "c"
+      end
     end
   end
 end