diff --git a/src/common/binaryheap.c b/src/common/binaryheap.c index 39a8243a6d..19e095f1fb 100644 --- a/src/common/binaryheap.c +++ b/src/common/binaryheap.c @@ -215,6 +215,35 @@ binaryheap_remove_first(binaryheap *heap) return result; } +/* + * binaryheap_remove_node + * + * Removes the nth (zero based) node from the heap. The caller must ensure + * that there are at least (n + 1) nodes in the heap. O(log n) worst case. + */ +void +binaryheap_remove_node(binaryheap *heap, int n) +{ + int cmp; + + Assert(!binaryheap_empty(heap) && heap->bh_has_heap_property); + Assert(n >= 0 && n < heap->bh_size); + + /* compare last node to the one that is being removed */ + cmp = heap->bh_compare(heap->bh_nodes[--heap->bh_size], + heap->bh_nodes[n], + heap->bh_arg); + + /* remove the last node, placing it in the vacated entry */ + heap->bh_nodes[n] = heap->bh_nodes[heap->bh_size]; + + /* sift as needed to preserve the heap property */ + if (cmp > 0) + sift_up(heap, n); + else if (cmp < 0) + sift_down(heap, n); +} + /* * binaryheap_replace_first * diff --git a/src/include/lib/binaryheap.h b/src/include/lib/binaryheap.h index 3647aeae65..9525dcaec4 100644 --- a/src/include/lib/binaryheap.h +++ b/src/include/lib/binaryheap.h @@ -59,8 +59,11 @@ extern void binaryheap_build(binaryheap *heap); extern void binaryheap_add(binaryheap *heap, bh_node_type d); extern bh_node_type binaryheap_first(binaryheap *heap); extern bh_node_type binaryheap_remove_first(binaryheap *heap); +extern void binaryheap_remove_node(binaryheap *heap, int n); extern void binaryheap_replace_first(binaryheap *heap, bh_node_type d); #define binaryheap_empty(h) ((h)->bh_size == 0) +#define binaryheap_size(h) ((h)->bh_size) +#define binaryheap_get_node(h, n) ((h)->bh_nodes[n]) #endif /* BINARYHEAP_H */