From f3ffb0da78314e40615192ff719e28e145870894 Mon Sep 17 00:00:00 2001 From: Tucker Evans Date: Tue, 21 Jul 2020 15:46:40 -0400 Subject: Add AVL checks to map remove --- collections/map/map.c | 25 ++++++++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/collections/map/map.c b/collections/map/map.c index 5fda2a5..015a3cb 100644 --- a/collections/map/map.c +++ b/collections/map/map.c @@ -462,9 +462,18 @@ void *key; tmp->parent = tmp->right = tmp->left = NULL; free(tmp); - *root = next; + if (!(*root = next) && !next_parent) + return ret; - //CHECK BALANCE + for (tmp = next_parent; tmp && tmp->parent; tmp = tmp->parent) { + bal = map_bal_factor(tmp); + if (bal <= -2 || bal >= 2) { + if (tmp == tmp->parent->left) { + map_rebal_tree(&(tmp->parent->left)); + } else if (tmp == tmp->parent->right) + map_rebal_tree(&(tmp->parent->right)); + } + } return ret; } @@ -472,10 +481,20 @@ void* map_remove(root, key) map *root; void *key; { + void *ret; + int bal; + if (!root) return NULL; - return map_remove_aux(&root->root, root->cmp, key); + ret = map_remove_aux(&root->root, root->cmp, key); + + bal = map_bal_factor(root->root); + if (bal <= -2 || bal >= 2) { + map_rebal_tree(&root->root); + } + + return ret; } void map_clear_aux(root) -- cgit v1.1