src/common/binaryheap.c

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