diff --git a/pyavltree.py b/pyavltree.py
index afa6614..bf3e50c 100755
--- a/pyavltree.py
+++ b/pyavltree.py
@@ -371,6 +371,9 @@ def remove_leaf (self, node):
         
     def remove_branch (self, node):
         parent = node.parent
+        leftChild = node.leftChild
+        rightChild = node.rightChild
+
         if (parent):
             if parent.leftChild == node:
                 parent.leftChild = node.rightChild or node.leftChild
@@ -386,10 +389,19 @@ def remove_branch (self, node):
         del node
         # rebalance
         node = parent
-        while (node):
-            if not node.balance() in [-1, 0, 1]:
-                self.rebalance(node)
-            node = node.parent
+
+        if node:
+            while (node):
+                if not node.balance() in [-1, 0, 1]:
+                    self.rebalance(node)
+                node = node.parent
+        else:
+            if leftChild:
+                self.rootNode = leftChild
+            else:
+                self.rootNode = rightChild
+
+            self.rootNode.parent = None
         
     def swap_with_successor_and_remove (self, node):
         successor = self.find_smallest(node.rightChild)