bug fixed, AVL trees now working properly

Author: Alejandro Fernández Suárez

include/avl.h

@@ -51,8 +51,9 @@ class AVL {
 
         void erase(const T &x) {
             if (!isEmpty()) {
-                tree = tree->erase(x);
-                numNodes--;
+                bool found = false;
+                tree = tree->erase(x, found);
+                if (found) numNodes--;
             }
         }
 
@@ -64,7 +65,7 @@ class AVL {
             }
         }
 
-    private:
+    public:
 
         struct Node {
             Node *left, *right;
@@ -93,51 +94,58 @@ class AVL {
                 return update();
             }
 
-            Node *erase(const T &x) {
+            Node *erase(const T &x, bool &found) {
                 if (value == x) {
+                    found = true;
                     if (left == 0 && right == 0) {
                         delete this;
                         return 0;
                     } else if (left == 0) {
+                        Node *aux = right;
                         *this = *right;
-                        delete right;
+                        delete aux;
                     } else if (right == 0) {
+                        Node *aux = left;
                         *this = *left;
-                        delete left;
+                        delete aux;
                     } else {
                         // Tracing path to rightmost leaf of the left subtree
                         std::stack<Node*> trace;
 
                         Node *current = left;
-                        while (current->right != 0) {
+                        while (current != 0) {
                             trace.push(current);
                             current = current->right;
                         }
 
+                        current = trace.top();
                         value = current->value;
                         Node *lsubtree = current->left;
                         delete current;
+                        trace.pop();
 
-                        if (trace.empty()) trace.push(left);
-
-                        trace.top()->right = lsubtree;
-
-                        do {
-                            trace.top()->update();
+                        if (trace.empty()) { left = lsubtree; }
+                        else {
+                            trace.top()->right = lsubtree;
                             trace.pop();
-                        } while (!trace.empty());
+                            while (!trace.empty()) {
+                                current = trace.top();
+                                current->right = current->right->update();
+                                trace.pop();
+                            }
+                        }
                     }
                     return update();
                 }
                 else if (x < value) {
                     if (left != 0) {
-                        left = left->erase(x);
+                        left = left->erase(x, found);
                         return update();
                     } else return this;
                 }
                 else {
                     if (right != 0) {
-                        right = right->erase(x);
+                        right = right->erase(x, found);
                         return update();
                     } else return this;
                 }

src/avl_demo.cpp

@@ -5,7 +5,7 @@ using namespace std;
 using namespace alg;
 
 const unsigned N = 4096*32;
-const unsigned N_ELEMS_TO_REMOVE = 4096*8;
+const unsigned N_ELEMS_TO_REMOVE = N-128; // Must be between 0 and N-1
 
 template <typename T> 
 void printTreeStatus(const AVL<T> &t) {
@@ -40,20 +40,27 @@ int main()
 	        cout << "ERROR: Value " << values[idx] << " was inserted and not found!" << endl;
     }
 
-    cout << "Now removing " << N_ELEMS_TO_REMOVE << " random elements for the tree... ";
+    cout << "Now removing a random element from the tree... ";
+    unsigned idx = rand() % N;
+    avl.erase(values[idx]);
+    cout << "Done" << endl;
+
+    printTreeStatus(avl);
+    
+    cout << "Now removing the root of the tree " << N_ELEMS_TO_REMOVE << " times... ";
     for (unsigned i = 0; i < N_ELEMS_TO_REMOVE; ++i) {
-        unsigned idx = rand() % N;
-        avl.erase(values[idx]);
+        avl.erase(avl.root());
     }
     cout << "Done" << endl;
 
     printTreeStatus(avl);
 
-    cout << "Do you want to see the GraphViz representation of the Tree (Y/n)? ";
+    // Outputting to cerr so the output can be redirected with ./avl_demo 2> <name>.gvz
+    cout << "Do you want to output the GraphViz representation of the tree to the cerr stream (Y/n)? ";
     char usrInput;
     cin >> usrInput;
-    if (usrInput == 'Y' || usrInput == 'y') avl.toGraphViz(cout, "AVL");
-
+    if (usrInput == 'Y' || usrInput == 'y') avl.toGraphViz(cerr, "AVL");
+    
     return 0;
 }