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
+ *         Generic K.Rege
  */
 public class AVLSearchTree<T extends Comparable<T>> extends SortedBinaryTree<T> {
     private boolean balanced(TreeNode<T> node) {
         // TODO Implement (6.4)
-        return true;
+        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() {
@@ -19,7 +26,7 @@ public class AVLSearchTree<T extends Comparable<T>> extends SortedBinaryTree<T>
     @Override
     protected int calcSize(TreeNode<T> p) {
         // TODO Implement (6.2)
-        return super.calcSize(p);
+        return super.calcSize(p) + p.values.size() - 1;
     }
 
     /**
@@ -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,8 +71,9 @@ 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.
+     * @param p       the node that roots the tree.
      * @return the new root.
      */
     private TreeNode<T> insertAt(TreeNode<T> p, T element) {
@@ -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) {

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

@@ -74,14 +74,19 @@ public class SortedBinaryTree<T extends Comparable<T>> implements Tree<T> {
     }
 
     protected int calcHeight(TreeNode<T> node) {
-        int result = 0;
+        if (node == null) {
+            return 0;
+        } else {
+            int result = 1;
 
-        if (node.left != null || node.right != null) {
-            result = Stream.of(node.left, node.right).filter(
-                (innerNode) -> innerNode != null).mapToInt((innerNode) -> calcHeight(innerNode)).max().getAsInt();
+            if (node.left != null || node.right != null) {
+                result += Stream.of(node.left, node.right).filter(
+                        (innerNode) -> innerNode != null).mapToInt((innerNode) -> calcHeight(innerNode)).max()
+                        .getAsInt();
+            }
+
+            return result;
         }
-
-        return result + 1;
     }
 
     public int height() {
@@ -89,14 +94,18 @@ public class SortedBinaryTree<T extends Comparable<T>> implements Tree<T> {
     }
 
     protected int calcSize(TreeNode<T> p) {
-        int result = 1;
+        int result = 0;
 
-        if (p.left != null) {
-            result += calcSize(p.left);
-        }
+        if (p != null) {
+            result++;
 
-        if (p.right != null) {
-            result += calcSize(p.right);
+            if (p.left != null) {
+                result += calcSize(p.left);
+            }
+
+            if (p.right != null) {
+                result += calcSize(p.right);
+            }
         }
 
         return result;