Re-balance tree after node removal

app/src/main/java/ch/nuth/zhaw/exbox/AVLSearchTree.java

@@ -3,13 +3,20 @@ package ch.nuth.zhaw.exbox;
 /**
  * Implements an AVL tree.
  * Note that all "matching" is based on the compareTo method.
+ * 
  * @author Mark Allen Weiss
  *         Generic K.Rege
  */
 public class AVLSearchTree<T extends Comparable<T>> extends SortedBinaryTree<T> {
     private boolean balanced(TreeNode<T> node) {
         // TODO Implement (6.4)
+        if (node == null) {
             return true;
+        } else if (balanced(node.left) && balanced(node.right)) {
+            return Math.abs(calcHeight(node.left) - calcHeight(node.right)) < 2;
+        } else {
+            return false;
+        }
     }
 
     public boolean balanced() {
@@ -31,6 +38,7 @@ public class AVLSearchTree<T extends Comparable<T>> extends SortedBinaryTree<T>
 
     /**
      * Insert into the tree; duplicates are ignored.
+     * 
      * @param element the item to insert.
      */
     public void add(T element) {
@@ -63,6 +71,7 @@ public class AVLSearchTree<T extends Comparable<T>> extends SortedBinaryTree<T>
 
     /**
      * Internal method to insert into a subtree.
+     * 
      * @param element the item to insert.
      * @param p       the node that roots the tree.
      * @return the new root.
@@ -88,6 +97,7 @@ public class AVLSearchTree<T extends Comparable<T>> extends SortedBinaryTree<T>
 
     // find node to replace
     private TreeNode<T> rep;
+
     private TreeNode<T> findRepAt(TreeNode<T> node) {
         if (node.right != null) {
             node.right = findRepAt(node.right);
@@ -129,13 +139,15 @@ public class AVLSearchTree<T extends Comparable<T>> extends SortedBinaryTree<T>
                 // search right
                 node.right = removeAt(node.right, x);
             }
+
             // TODO Implement (6.5)
-            return node;
+            return balance(node);
         }
     }
 
     /**
      * Remove from the tree. Nothing is done if x is not found.
+     * 
      * @param x the item to remove.
      */
     public T remove(T x) {