1 files changed, 16 insertions, 16 deletions

diff --git a/tree.c b/tree.c
index 50f9548..3c6a65c 100644
--- a/tree.c
+++ b/tree.c
@@ -132,27 +132,27 @@ tree_remove(struct tree_node *tree, T key)
                 tree->left = tree_remove(tree->left, key);
         else if ( key > tree->key )
                 tree->right = tree_remove(tree->right, key);
-        else { /* key == tree->key */
-                /* TODO: Handle Collision */
+        else {                          /* key == tree->key */
+                if ( --tree->count == 0 ) { /* Handle Collision */
+                        struct tree_node *temp = tree;
 
-                struct tree_node *temp = tree;
+                        if ( tree->left == NULL ) {
+                                tree = tree->right;
+                        }
+                        else if ( tree->right == NULL ) {
+                                tree = tree->left;
+                        }
+                        else {
+                                struct tree_node *min = tree_detach_min(&tree->right);
 
-                if ( tree->left == NULL ) {
-                        tree = tree->right;
-                }
-                else if ( tree->right == NULL ) {
-                        tree = tree->left;
-                }
-                else {
-                        struct tree_node *min = tree_detach_min(&tree->right);
+                                min->left  = tree->left;
+                                min->right = tree->right;
 
-                        min->left  = tree->left;
-                        min->right = tree->right;
+                                tree = min;
+                        }
 
-                        tree = min;
+                        free(temp);
                 }
-
-                free(temp);
         }
         return tree;
 }