Check the balance of an arbitrary tree.

1 files changed, 50 insertions, 0 deletions

diff --git a/tree.cpp b/tree.cpp
index 690c964..4172230 100644
--- a/tree.cpp
+++ b/tree.cpp
@@ -82,6 +82,50 @@ void insert_value(node* root, long value)
     }
 }
 
+bool count_depth(const node* nd, int& depth)
+{
+    assert(nd);
+
+    int d_left = 0;
+    int d_right = 0;
+    if (nd->left)
+    {
+        if (!count_depth(nd->left, d_left))
+            return false;
+    }
+
+    if (nd->right)
+    {
+        if (!count_depth(nd->right, d_right))
+            return false;
+    }
+
+    cout << "(" << nd->value << ") left:" << d_left << " right:" << d_right << endl;
+
+    if (d_left > d_right && d_left-d_right > 1)
+        return false;
+
+    if (d_right > d_left && d_right-d_left > 1)
+        return false;
+
+    depth += 1 + (d_left > d_right ? d_left : d_right);
+    return true;
+}
+
+void check_balance(const node* root)
+{
+    cout << "checking the balance of the tree..." << endl;
+    if (!root)
+        return;
+
+    int depth = -1;
+    bool res = count_depth(root, depth);
+    if (res)
+        cout << "tree is balanced (depth: " << depth << ")" << endl;
+    else
+        cout << "tree is not balanced" << endl;
+}
+
 int main()
 {
     node* root = NULL;
@@ -99,6 +143,12 @@ int main()
             break;
         }
 
+        if (val == "c" || val == "check")
+        {
+            check_balance(root);
+            continue;
+        }
+
         // End position on successful conversion.
         const char* end_pos = val.c_str() + val.size();