diff --git a/src/tree/balancing.c b/src/tree/balancing.c
index b7466ea..9f9d896 100644
--- a/src/tree/balancing.c
+++ b/src/tree/balancing.c
@@ -3,22 +3,35 @@
 vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
   vieter_tree_node *parent = node->parent;
   vieter_tree_node *grand_parent = parent->parent;
-  vieter_tree_node *grand_grand_parent = grand_parent->parent;
-  vieter_tree_node *root, *left, *right;
+  vieter_tree_node *root = grand_parent;
+  vieter_tree_node *left, *right;
   vieter_tree_node *children[4];
 
+  uint64_t key_root = root->key;
+  void *data_root = root->data;
+
   if (node->key < parent->key && parent->key < grand_parent->key) {
-    root = parent;
+    root->key = parent->key;
+    root->data = parent->data;
+
+    parent->key = key_root;
+    parent->data = data_root;
+
     left = node;
-    right = grand_parent;
+    right = parent;
 
     children[0] = node->left;
     children[1] = node->right;
     children[2] = parent->right;
     children[3] = grand_parent->right;
   } else if (node->key < parent->key && parent->key > grand_parent->key) {
-    root = node;
-    left = grand_parent;
+    root->key = node->key;
+    root->data = node->data;
+
+    node->key = key_root;
+    node->data = data_root;
+
+    left = node;
     right = parent;
 
     children[0] = grand_parent->left;
@@ -26,17 +39,27 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
     children[2] = node->right;
     children[3] = parent->right;
   } else if (node->key > parent->key && parent->key < grand_parent->key) {
-    root = node;
+    root->key = node->key;
+    root->data = node->data;
+
+    node->key = key_root;
+    node->data = data_root;
+
     left = parent;
-    right = grand_parent;
+    right = node;
 
     children[0] = parent->left;
     children[1] = node->left;
     children[2] = node->right;
     children[3] = grand_parent->right;
   } else {
-    root = parent;
-    left = grand_parent;
+    root->key = parent->key;
+    root->data = parent->data;
+
+    parent->key = key_root;
+    parent->data = data_root;
+
+    left = parent;
     right = node;
 
     children[0] = grand_parent->left;
@@ -72,12 +95,6 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
   left->parent = root;
   right->parent = root;
 
-  if (grand_grand_parent != NULL) {
-    vieter_tree_node_add_child(grand_grand_parent, root->key, root);
-  } else {
-    root->parent = NULL;
-  }
-
   return root;
 }
 
diff --git a/src/tree/node.c b/src/tree/node.c
index d6c2b6d..0d74773 100644
--- a/src/tree/node.c
+++ b/src/tree/node.c
@@ -24,12 +24,18 @@ void vieter_tree_node_add_child(vieter_tree_node *parent, uint64_t key,
   }
 }
 
-void vieter_tree_node_replace(vieter_tree_node *to_replace,
-                              vieter_tree_node *replacement) {
+void vieter_tree_node_replace_with_child(vieter_tree_node *to_replace,
+                                         vieter_tree_node *replacement) {
   to_replace->key = replacement->key;
   to_replace->data = replacement->data;
   to_replace->left = replacement->left;
   to_replace->right = replacement->right;
+
+  if (to_replace->left != NULL)
+    to_replace->left->parent = to_replace;
+
+  if (to_replace->right != NULL)
+    to_replace->right->parent = to_replace;
 }
 
 vieter_tree_error vieter_tree_node_insert(vieter_tree_node *root, uint64_t key,
@@ -120,7 +126,7 @@ vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root,
       vieter_tree_node_add_child(target->parent, child->key, child);
       vieter_tree_node_free(target);
     } else {
-      vieter_tree_node_replace(target, child);
+      vieter_tree_node_replace_with_child(target, child);
       vieter_tree_node_free(child);
     }
 
diff --git a/test/tree/test_binary_tree.c b/test/tree/test_binary_tree.c
index 7e85ebe..aa7fc0f 100644
--- a/test/tree/test_binary_tree.c
+++ b/test/tree/test_binary_tree.c
@@ -15,14 +15,14 @@ void test_init() {
 void test_insert() {
     vieter_tree *tree = vieter_tree_init();
 
-    for (uint64_t i = 0; i < 25; i++) {
+    for (uint64_t i = 0; i < 250; i++) {
         TEST_CHECK(vieter_tree_insert(tree, i, NULL) == vieter_tree_ok);
         TEST_SIZE(tree, i + 1);
     }
 
     void *out;
 
-    for (uint64_t i = 0; i < 25; i++) {
+    for (uint64_t i = 0; i < 250; i++) {
         TEST_CHECK(vieter_tree_search(&out, tree, i) == vieter_tree_ok);
         TEST_CHECK(vieter_tree_insert(tree, i, NULL) == vieter_tree_already_present);
         TEST_CHECK(vieter_tree_search(&out, tree, i) == vieter_tree_ok);
@@ -34,17 +34,17 @@ void test_insert() {
 void test_remove() {
     vieter_tree *tree = vieter_tree_init();
 
-    for (uint64_t i = 0; i < 25; i++) {
+    for (uint64_t i = 0; i < 250; i++) {
         TEST_CHECK(vieter_tree_insert(tree, i, NULL) == vieter_tree_ok);
     }
 
     void *out;
 
-    for (uint64_t i = 0; i < 25; i++) {
+    for (uint64_t i = 0; i < 250; i++) {
         TEST_CHECK(vieter_tree_search(&out, tree, i) == vieter_tree_ok);
         TEST_CHECK(vieter_tree_remove(&out, tree, i) == vieter_tree_ok);
         TEST_CHECK(vieter_tree_search(&out, tree, i) == vieter_tree_not_present);
-        TEST_SIZE(tree, 25 - i - 1);
+        TEST_SIZE(tree, 250 - i - 1);
     }
 
     vieter_tree_free(tree);