aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--src/common/binaryheap.c29
-rw-r--r--src/include/lib/binaryheap.h3
2 files changed, 32 insertions, 0 deletions
diff --git a/src/common/binaryheap.c b/src/common/binaryheap.c
index 39a8243a6d5..19e095f1fb1 100644
--- a/src/common/binaryheap.c
+++ b/src/common/binaryheap.c
@@ -216,6 +216,35 @@ binaryheap_remove_first(binaryheap *heap)
}
/*
+ * 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
*
* Replace the topmost element of a non-empty heap, preserving the heap
diff --git a/src/include/lib/binaryheap.h b/src/include/lib/binaryheap.h
index 3647aeae657..9525dcaec44 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 */