diff --git a/src/heap/vieter_heap.c b/src/heap/vieter_heap.c index 8e036b6..dee109c 100644 --- a/src/heap/vieter_heap.c +++ b/src/heap/vieter_heap.c @@ -57,13 +57,15 @@ vieter_heap_error vieter_heap_pop(void **out, vieter_heap *heap) { } if (heap->tree->order == 0) { - *out = heap->tree->root->data; - vieter_heap_tree_free(heap->tree); - heap->tree = NULL; + *out = heap->tree->root->data; + vieter_heap_tree_free(heap->tree); + heap->tree = NULL; - return vieter_heap_ok; + return vieter_heap_ok; } + vieter_heap_tree_pop(out, heap->tree); + return vieter_heap_ok; } diff --git a/src/heap/vieter_heap_tree.c b/src/heap/vieter_heap_tree.c index 8dcc29b..9a170e3 100644 --- a/src/heap/vieter_heap_tree.c +++ b/src/heap/vieter_heap_tree.c @@ -123,7 +123,7 @@ vieter_heap_tree *vieter_heap_tree_merge(vieter_heap_tree *tree_a, tree = next_tree; } else if (target->order > tree->order) { - next_tree = tree->next; + next_tree = tree->next; if (previous_target == NULL) { previous_target = tree; @@ -152,6 +152,50 @@ vieter_heap_tree *vieter_heap_tree_merge(vieter_heap_tree *tree_a, return out; } -vieter_heap_tree *vieter_heap_tree_pop(vieter_heap_tree *tree) { +vieter_heap_tree *vieter_heap_tree_pop(void **out, vieter_heap_tree *tree) { + vieter_heap_tree *tree_before_smallest = NULL; + vieter_heap_tree *previous_tree = NULL; + vieter_heap_tree *tree_out = tree; + uint64_t smallest_key = tree->root->key; + + while (tree->next != NULL) { + previous_tree = tree; + tree = tree->next; + + if (tree->root->key < smallest_key) { + smallest_key = tree->root->key; + *out = tree->root->data; + tree_before_smallest = previous_tree; + } + } + + if (tree_before_smallest != NULL) { + previous_tree->next = tree->next; + } else { + tree_out = tree_out->next; + } + + if (tree->order == 0) { + vieter_heap_tree_free(tree); + + return NULL; + } + + uint8_t old_order = tree->order; + + vieter_heap_node *node = tree->root->largest_order; + free(tree->root); + + previous_tree = vieter_heap_tree_init(node, NULL, old_order - 1); + + uint8_t i = 2; + while (node->next_largest_order != NULL) { + node = node->next_largest_order; + previous_tree = vieter_heap_tree_init(node, previous_tree, old_order - i); + + i++; + } + + return vieter_heap_tree_merge(tree_out, previous_tree); } diff --git a/src/heap/vieter_heap_tree.h b/src/heap/vieter_heap_tree.h index 117e1a9..10c6256 100644 --- a/src/heap/vieter_heap_tree.h +++ b/src/heap/vieter_heap_tree.h @@ -33,6 +33,6 @@ vieter_heap_tree *vieter_heap_tree_merge(vieter_heap_tree *tree_a, vieter_heap_t vieter_heap_tree *vieter_heap_tree_merge_same_order(vieter_heap_tree *tree_a, vieter_heap_tree *tree_b); -vieter_heap_tree *vieter_heap_tree_pop(vieter_heap_tree *tree); +vieter_heap_tree *vieter_heap_tree_pop(void **out, vieter_heap_tree *tree); #endif diff --git a/test/heap/test_heap.c b/test/heap/test_heap.c index 38d9095..dbc26b8 100644 --- a/test/heap/test_heap.c +++ b/test/heap/test_heap.c @@ -59,7 +59,7 @@ void test_pop() { for (uint64_t i = 50; i > 0; i--) { vieter_heap_insert(heap, i, (void *)i); - TEST_SIZE(heap, (uint64_t)50 - i); + TEST_SIZE(heap, (uint64_t)51 - i); TEST_CHECK(vieter_heap_peek(&data, heap) == vieter_heap_ok); TEST_CHECK(data == (void*)i); @@ -80,5 +80,6 @@ TEST_LIST = { {"init", test_init}, {"merge same order", test_merge_same_order}, {"insert", test_insert}, + {"pop", test_pop}, {NULL, NULL} };