granicus.if.org Git - postgresql/blob - src/backend/lib/binaryheap.c

   1 /*-------------------------------------------------------------------------
   2  *
   3  * binaryheap.c
   4  *        A simple binary heap implementaion
   5  *
   6  * Portions Copyright (c) 2012-2013, PostgreSQL Global Development Group
   7  *
   8  * IDENTIFICATION
   9  *        src/backend/lib/binaryheap.c
  10  *
  11  *-------------------------------------------------------------------------
  12  */
  13
  14 #include "postgres.h"
  15
  16 #include <math.h>
  17
  18 #include "lib/binaryheap.h"
  19
  20 static void sift_down(binaryheap *heap, int node_off);
  21 static void sift_up(binaryheap *heap, int node_off);
  22 static inline void swap_nodes(binaryheap *heap, int a, int b);
  23
  24 /*
  25  * binaryheap_allocate
  26  *
  27  * Returns a pointer to a newly-allocated heap that has the capacity to
  28  * store the given number of nodes, with the heap property defined by
  29  * the given comparator function, which will be invoked with the additional
  30  * argument specified by 'arg'.
  31  */
  32 binaryheap *
  33 binaryheap_allocate(int capacity, binaryheap_comparator compare, void *arg)
  34 {
  35         int                     sz;
  36         binaryheap *heap;
  37
  38         sz = offsetof(binaryheap, bh_nodes) +sizeof(Datum) * capacity;
  39         heap = (binaryheap *) palloc(sz);
  40         heap->bh_space = capacity;
  41         heap->bh_compare = compare;
  42         heap->bh_arg = arg;
  43
  44         heap->bh_size = 0;
  45         heap->bh_has_heap_property = true;
  46
  47         return heap;
  48 }
  49
  50 /*
  51  * binaryheap_reset
  52  *
  53  * Resets the heap to an empty state, losing its data content but not the
  54  * parameters passed at allocation.
  55  */
  56 void
  57 binaryheap_reset(binaryheap *heap)
  58 {
  59         heap->bh_size = 0;
  60         heap->bh_has_heap_property = true;
  61 }
  62
  63 /*
  64  * binaryheap_free
  65  *
  66  * Releases memory used by the given binaryheap.
  67  */
  68 void
  69 binaryheap_free(binaryheap *heap)
  70 {
  71         pfree(heap);
  72 }
  73
  74 /*
  75  * These utility functions return the offset of the left child, right
  76  * child, and parent of the node at the given index, respectively.
  77  *
  78  * The heap is represented as an array of nodes, with the root node
  79  * stored at index 0. The left child of node i is at index 2*i+1, and
  80  * the right child at 2*i+2. The parent of node i is at index (i-1)/2.
  81  */
  82
  83 static inline int
  84 left_offset(int i)
  85 {
  86         return 2 * i + 1;
  87 }
  88
  89 static inline int
  90 right_offset(int i)
  91 {
  92         return 2 * i + 2;
  93 }
  94
  95 static inline int
  96 parent_offset(int i)
  97 {
  98         return (i - 1) / 2;
  99 }
 100
 101 /*
 102  * binaryheap_add_unordered
 103  *
 104  * Adds the given datum to the end of the heap's list of nodes in O(1) without
 105  * preserving the heap property. This is a convenience to add elements quickly
 106  * to a new heap. To obtain a valid heap, one must call binaryheap_build()
 107  * afterwards.
 108  */
 109 void
 110 binaryheap_add_unordered(binaryheap *heap, Datum d)
 111 {
 112         if (heap->bh_size >= heap->bh_space)
 113                 elog(ERROR, "out of binary heap slots");
 114         heap->bh_has_heap_property = false;
 115         heap->bh_nodes[heap->bh_size] = d;
 116         heap->bh_size++;
 117 }
 118
 119 /*
 120  * binaryheap_build
 121  *
 122  * Assembles a valid heap in O(n) from the nodes added by
 123  * binaryheap_add_unordered(). Not needed otherwise.
 124  */
 125 void
 126 binaryheap_build(binaryheap *heap)
 127 {
 128         int                     i;
 129
 130         for (i = parent_offset(heap->bh_size - 1); i >= 0; i--)
 131                 sift_down(heap, i);
 132         heap->bh_has_heap_property = true;
 133 }
 134
 135 /*
 136  * binaryheap_add
 137  *
 138  * Adds the given datum to the heap in O(log n) time, while preserving
 139  * the heap property.
 140  */
 141 void
 142 binaryheap_add(binaryheap *heap, Datum d)
 143 {
 144         if (heap->bh_size >= heap->bh_space)
 145                 elog(ERROR, "out of binary heap slots");
 146         heap->bh_nodes[heap->bh_size] = d;
 147         heap->bh_size++;
 148         sift_up(heap, heap->bh_size - 1);
 149 }
 150
 151 /*
 152  * binaryheap_first
 153  *
 154  * Returns a pointer to the first (root, topmost) node in the heap
 155  * without modifying the heap. The caller must ensure that this
 156  * routine is not used on an empty heap. Always O(1).
 157  */
 158 Datum
 159 binaryheap_first(binaryheap *heap)
 160 {
 161         Assert(!binaryheap_empty(heap) && heap->bh_has_heap_property);
 162         return heap->bh_nodes[0];
 163 }
 164
 165 /*
 166  * binaryheap_remove_first
 167  *
 168  * Removes the first (root, topmost) node in the heap and returns a
 169  * pointer to it after rebalancing the heap. The caller must ensure
 170  * that this routine is not used on an empty heap. O(log n) worst
 171  * case.
 172  */
 173 Datum
 174 binaryheap_remove_first(binaryheap *heap)
 175 {
 176         Assert(!binaryheap_empty(heap) && heap->bh_has_heap_property);
 177
 178         if (heap->bh_size == 1)
 179         {
 180                 heap->bh_size--;
 181                 return heap->bh_nodes[0];
 182         }
 183
 184         /*
 185          * Swap the root and last nodes, decrease the size of the heap (i.e.
 186          * remove the former root node) and sift the new root node down to its
 187          * correct position.
 188          */
 189         swap_nodes(heap, 0, heap->bh_size - 1);
 190         heap->bh_size--;
 191         sift_down(heap, 0);
 192
 193         return heap->bh_nodes[heap->bh_size];
 194 }
 195
 196 /*
 197  * binaryheap_replace_first
 198  *
 199  * Replace the topmost element of a non-empty heap, preserving the heap
 200  * property.  O(1) in the best case, or O(log n) if it must fall back to
 201  * sifting the new node down.
 202  */
 203 void
 204 binaryheap_replace_first(binaryheap *heap, Datum d)
 205 {
 206         Assert(!binaryheap_empty(heap) && heap->bh_has_heap_property);
 207
 208         heap->bh_nodes[0] = d;
 209
 210         if (heap->bh_size > 1)
 211                 sift_down(heap, 0);
 212 }
 213
 214 /*
 215  * Swap the contents of two nodes.
 216  */
 217 static inline void
 218 swap_nodes(binaryheap *heap, int a, int b)
 219 {
 220         Datum           swap;
 221
 222         swap = heap->bh_nodes[a];
 223         heap->bh_nodes[a] = heap->bh_nodes[b];
 224         heap->bh_nodes[b] = swap;
 225 }
 226
 227 /*
 228  * Sift a node up to the highest position it can hold according to the
 229  * comparator.
 230  */
 231 static void
 232 sift_up(binaryheap *heap, int node_off)
 233 {
 234         while (node_off != 0)
 235         {
 236                 int                     cmp;
 237                 int                     parent_off;
 238
 239                 /*
 240                  * If this node is smaller than its parent, the heap condition is
 241                  * satisfied, and we're done.
 242                  */
 243                 parent_off = parent_offset(node_off);
 244                 cmp = heap->bh_compare(heap->bh_nodes[node_off],
 245                                                            heap->bh_nodes[parent_off],
 246                                                            heap->bh_arg);
 247                 if (cmp <= 0)
 248                         break;
 249
 250                 /*
 251                  * Otherwise, swap the node and its parent and go on to check the
 252                  * node's new parent.
 253                  */
 254                 swap_nodes(heap, node_off, parent_off);
 255                 node_off = parent_off;
 256         }
 257 }
 258
 259 /*
 260  * Sift a node down from its current position to satisfy the heap
 261  * property.
 262  */
 263 static void
 264 sift_down(binaryheap *heap, int node_off)
 265 {
 266         while (true)
 267         {
 268                 int                     left_off = left_offset(node_off);
 269                 int                     right_off = right_offset(node_off);
 270                 int                     swap_off = 0;
 271
 272                 /* Is the left child larger than the parent? */
 273                 if (left_off < heap->bh_size &&
 274                         heap->bh_compare(heap->bh_nodes[node_off],
 275                                                          heap->bh_nodes[left_off],
 276                                                          heap->bh_arg) < 0)
 277                         swap_off = left_off;
 278
 279                 /* Is the right child larger than the parent? */
 280                 if (right_off < heap->bh_size &&
 281                         heap->bh_compare(heap->bh_nodes[node_off],
 282                                                          heap->bh_nodes[right_off],
 283                                                          heap->bh_arg) < 0)
 284                 {
 285                         /* swap with the larger child */
 286                         if (!swap_off ||
 287                                 heap->bh_compare(heap->bh_nodes[left_off],
 288                                                                  heap->bh_nodes[right_off],
 289                                                                  heap->bh_arg) < 0)
 290                                 swap_off = right_off;
 291                 }
 292
 293                 /*
 294                  * If we didn't find anything to swap, the heap condition is
 295                  * satisfied, and we're done.
 296                  */
 297                 if (!swap_off)
 298                         break;
 299
 300                 /*
 301                  * Otherwise, swap the node with the child that violates the heap
 302                  * property; then go on to check its children.
 303                  */
 304                 swap_nodes(heap, swap_off, node_off);
 305                 node_off = swap_off;
 306         }
 307 }