diff --git a/src/tree/vieter_tree_balancing.c b/src/tree/vieter_tree_balancing.c
index 6a58179..b160563 100644
--- a/src/tree/vieter_tree_balancing.c
+++ b/src/tree/vieter_tree_balancing.c
@@ -33,15 +33,15 @@ 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 *root = grand_parent;
-  vieter_tree_node *left, *right;
-  vieter_tree_node *children[4];
+  vieter_tree_node *children[2];
+  vieter_tree_node *grand_children[4];
 
   uint64_t key_root = root->key;
   void *data_root = root->data;
 
   if (vieter_tree_node_get(node, vieter_tree_node_right)) {
-    left = parent;
-    right = node;
+    children[0] = parent;
+    children[1] = node;
 
     if (vieter_tree_node_get(parent, vieter_tree_node_right)) {
       root->key = parent->key;
@@ -50,10 +50,10 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
       parent->key = key_root;
       parent->data = data_root;
 
-      children[0] = grand_parent->children[0];
-      children[1] = parent->children[0];
-      children[2] = node->children[0];
-      children[3] = node->children[1];
+      grand_children[0] = grand_parent->children[0];
+      grand_children[1] = parent->children[0];
+      grand_children[2] = node->children[0];
+      grand_children[3] = node->children[1];
     } else {
       root->key = node->key;
       root->data = node->data;
@@ -61,14 +61,14 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
       node->key = key_root;
       node->data = data_root;
 
-      children[0] = parent->children[0];
-      children[1] = node->children[0];
-      children[2] = node->children[1];
-      children[3] = grand_parent->children[1];
+      grand_children[0] = parent->children[0];
+      grand_children[1] = node->children[0];
+      grand_children[2] = node->children[1];
+      grand_children[3] = grand_parent->children[1];
     }
   } else {
-    left = node;
-    right = parent;
+    children[0] = node;
+    children[1] = parent;
 
     if (vieter_tree_node_get(parent, vieter_tree_node_right)) {
       root->key = node->key;
@@ -77,10 +77,10 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
       node->key = key_root;
       node->data = data_root;
 
-      children[0] = grand_parent->children[0];
-      children[1] = node->children[0];
-      children[2] = node->children[1];
-      children[3] = parent->children[1];
+      grand_children[0] = grand_parent->children[0];
+      grand_children[1] = node->children[0];
+      grand_children[2] = node->children[1];
+      grand_children[3] = parent->children[1];
     } else {
       root->key = parent->key;
       root->data = parent->data;
@@ -88,10 +88,10 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
       parent->key = key_root;
       parent->data = data_root;
 
-      children[0] = node->children[0];
-      children[1] = node->children[1];
-      children[2] = parent->children[1];
-      children[3] = grand_parent->children[1];
+      grand_children[0] = node->children[0];
+      grand_children[1] = node->children[1];
+      grand_children[2] = parent->children[1];
+      grand_children[3] = grand_parent->children[1];
     }
   }
 
@@ -99,41 +99,14 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
       root, vieter_tree_node_right,
       vieter_tree_node_get(grand_parent, vieter_tree_node_right));
 
-  vieter_tree_node_set(left, vieter_tree_node_black, true);
-  vieter_tree_node_set(left, vieter_tree_node_right, false);
-  vieter_tree_node_set(right, vieter_tree_node_black, true);
-  vieter_tree_node_set(left, vieter_tree_node_right, true);
+  vieter_tree_node_set(children[0], vieter_tree_node_black, true);
+  vieter_tree_node_set(children[1], vieter_tree_node_black, true);
 
   vieter_tree_node_set(root, vieter_tree_node_black, false);
 
-  memcpy(&left->children, children, 2 * sizeof(vieter_tree_node *));
-
-  if (children[0] != NULL) {
-    children[0]->parent = left;
-    vieter_tree_node_set(children[0], vieter_tree_node_right, false);
-  }
-
-  if (children[1] != NULL) {
-    children[1]->parent = left;
-    vieter_tree_node_set(children[0], vieter_tree_node_right, true);
-  }
-
-  memcpy(&right->children, children + 2, 2 * sizeof(vieter_tree_node *));
-
-  if (children[2] != NULL) {
-    children[2]->parent = right;
-    vieter_tree_node_set(children[2], vieter_tree_node_right, false);
-  }
-
-  if (children[3] != NULL) {
-    children[3]->parent = right;
-    vieter_tree_node_set(children[2], vieter_tree_node_right, true);
-  }
-
-  root->children[0] = left;
-  root->children[1] = right;
-  left->parent = root;
-  right->parent = root;
+  vieter_tree_node_set_children(children[0], grand_children);
+  vieter_tree_node_set_children(children[1], grand_children + 2);
+  vieter_tree_node_set_children(root, children);
 
   return root;
 }
diff --git a/src/tree/vieter_tree_node.c b/src/tree/vieter_tree_node.c
index 853281a..2f45b82 100644
--- a/src/tree/vieter_tree_node.c
+++ b/src/tree/vieter_tree_node.c
@@ -24,18 +24,23 @@ void vieter_tree_node_add_child(vieter_tree_node *parent, uint64_t key,
   }
 }
 
+void vieter_tree_node_set_children(vieter_tree_node *parent,
+                                   vieter_tree_node **children) {
+  memcpy(parent->children, children, 2 * sizeof(vieter_tree_node *));
+
+  for (uint8_t i = 0; i < 2; i++) {
+    if (parent->children[i] != NULL) {
+      parent->children[i]->parent = parent;
+      vieter_tree_node_set(parent->children[i], vieter_tree_node_right, i);
+    }
+  }
+}
+
 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;
-  memcpy(to_replace->children, replacement->children,
-         2 * sizeof(vieter_tree_node *));
-
-  if (to_replace->children[0] != NULL)
-    to_replace->children[0]->parent = to_replace;
-
-  if (to_replace->children[1] != NULL)
-    to_replace->children[1]->parent = to_replace;
+  vieter_tree_node_set_children(to_replace, replacement->children);
 }
 
 vieter_tree_error vieter_tree_node_insert(vieter_tree_node *root, uint64_t key,
diff --git a/src/tree/vieter_tree_node.h b/src/tree/vieter_tree_node.h
index 4012d1d..1ebb15a 100644
--- a/src/tree/vieter_tree_node.h
+++ b/src/tree/vieter_tree_node.h
@@ -42,6 +42,7 @@ bool vieter_tree_node_get(vieter_tree_node *node, vieter_tree_node_flag flag);
 
 void vieter_tree_node_add_child(vieter_tree_node *parent, uint64_t key,
                                 vieter_tree_node *child);
+void vieter_tree_node_set_children(vieter_tree_node *parent, vieter_tree_node **children);
 
 bool vieter_tree_node_validate(vieter_tree_node *node, uint64_t passed_black_nodes, uint64_t expected_black_nodes);