include/c++/util/heap.h

   1 // util/heap.h -- A heap-based priority queue implementation
   2 //
   3 // This software is copyright (c) 2006 Scott Wood <scott@buserror.net>.
   4 //
   5 // This software is provided 'as-is', without any express or implied warranty.
   6 // In no event will the authors or contributors be held liable for any damages
   7 // arising from the use of this software.
   8 //
   9 // Permission is hereby granted to everyone, free of charge, to use, copy,
  10 // modify, prepare derivative works of, publish, distribute, perform,
  11 // sublicense, and/or sell copies of the Software, provided that the above
  12 // copyright notice and disclaimer of warranty be included in all copies or
  13 // substantial portions of this software.
  14
  15 #ifndef _UTIL_HEAP_H
  16 #define _UTIL_HEAP_H
  17
  18 #include <assert.h>
  19 #include <stddef.h>
  20 #include <stdint.h>
  21 #include <lowlevel/bitops.h>
  22
  23 namespace Util {
  24         // T must support operator < for key ordering, and
  25         // have a Heap<T>::Node member called heap_node.
  26
  27         template <typename T>
  28         class Heap {
  29         public:
  30                 struct Node {
  31                         Node *parent, *left, *right;
  32
  33                         // A pointer to the parent's left or right pointer
  34                         // (whichever side this node is on).  This is also
  35                         // used for sanity checks (it is NULL for nodes not
  36                         // on a queue).
  37
  38                         Node **parentlink;
  39
  40                         Node()
  41                         {
  42                                 parent = left = right = NULL;
  43                                 parentlink = NULL;
  44                         }
  45
  46                         bool is_on_heap()
  47                         {
  48                                 return parentlink != NULL;
  49                         }
  50                 };
  51
  52         private:
  53                 Node *top;
  54                 size_t num_nodes;
  55
  56                 T *node_to_type(Node *n)
  57                 {
  58                         return reinterpret_cast<T *>(reinterpret_cast<uintptr_t>(n) -
  59                                                      offsetof(T, heap_node));
  60                 }
  61
  62                 // Swap n with its parent
  63                 void swap(Node *n)
  64                 {
  65                         Node *p = n->parent;
  66                         *p->parentlink = n;
  67
  68                         n->parent = p->parent;
  69                         p->parent = n;
  70
  71                         Node **plink = p->parentlink;
  72
  73                         if (n->parentlink == &p->left) {
  74                                 assert(p->left == n);
  75
  76                                 Node *tmp = p->right;
  77                                 p->right = n->right;
  78                                 n->right = tmp;
  79
  80                                 p->left = n->left;
  81                                 n->left = p;
  82
  83                                 p->parentlink = &n->left;
  84
  85                                 if (n->right) {
  86                                         n->right->parent = n;
  87                                         n->right->parentlink = &n->right;
  88                                 }
  89                         } else {
  90                                 assert(n->parentlink == &p->right);
  91                                 assert(p->right == n);
  92
  93                                 Node *tmp = p->left;
  94                                 p->left = n->left;
  95                                 n->left = tmp;
  96
  97                                 p->right = n->right;
  98                                 n->right = p;
  99
 100                                 p->parentlink = &n->right;
 101
 102                                 if (n->left) {
 103                                         n->left->parent = n;
 104                                         n->left->parentlink = &n->left;
 105                                 }
 106                         }
 107
 108                         if (p->left) {
 109                                 p->left->parent = p;
 110                                 p->left->parentlink = &p->left;
 111
 112                                 if (p->right) {
 113                                         p->right->parent = p;
 114                                         p->right->parentlink = &p->right;
 115                                 }
 116                         } else {
 117                                 assert(!p->right);
 118                         }
 119
 120                         n->parentlink = plink;
 121                 }
 122
 123                 // Restore heap condition by moving n towards the root,
 124                 // if required.
 125
 126                 void up(Node *n)
 127                 {
 128                         T *n_type = node_to_type(n);
 129
 130                         assert(!n->left || !(*node_to_type(n->left) < *n_type));
 131                         assert(!n->right || !(*node_to_type(n->right) < *n_type));
 132
 133                         while (n->parent && *n_type < *node_to_type(n->parent)) {
 134                                 assert(!n->left || !(*node_to_type(n->left) < *n_type));
 135                                 assert(!n->right || !(*node_to_type(n->right) < *n_type));
 136
 137                                 swap(n);
 138                         }
 139
 140                         assert(!n->left || !(*node_to_type(n->left) < *n_type));
 141                         assert(!n->right || !(*node_to_type(n->right) < *n_type));
 142                 }
 143
 144                 // Restore heap condition by moving n towards the leaves,
 145                 // if required.
 146
 147                 void down(Node *n)
 148                 {
 149                         T *n_type = node_to_type(n);
 150
 151                         assert(!n->parent || !(*n_type < *node_to_type(n->parent)));
 152
 153                         // If n has a right, it'll have a left.
 154                         while (n->left) {
 155                                 assert(!n->parent || !(*n_type < *node_to_type(n->parent)));
 156                                 Node *child = n->left;
 157
 158                                 if (n->right &&
 159                                     *node_to_type(n->right) < *node_to_type(n->left))
 160                                         child = n->right;
 161
 162                                 if (*node_to_type(child) < *n_type)
 163                                         swap(child);
 164                                 else
 165                                         break;
 166                         }
 167
 168                         assert(!n->parent || !(*n_type < *node_to_type(n->parent)));
 169                 }
 170
 171                 Node *get_by_index(size_t num, Node **parent, Node ***link)
 172                 {
 173                         assert(num_nodes == 0 || top);
 174                         assert(num <= num_nodes);
 175
 176                         // It's easier to use 1-based indices, so that the high set bit
 177                         // perfectly indicates which level the node is on.
 178
 179                         num += 1;
 180
 181                         Node *n = NULL;
 182                         size_t bit = 1 << ll_get_order_round_down(num);
 183                         *parent = NULL;
 184                         Node **nextlink = &top;
 185
 186                         do {
 187                                 *link = nextlink;
 188                                 *parent = n;
 189                                 n = *nextlink;
 190
 191                                 bit /= 2;
 192
 193                                 if (num & bit)
 194                                         nextlink = &n->right;
 195                                 else
 196                                         nextlink = &n->left;
 197                         } while (bit);
 198
 199                         return n;
 200                 }
 201
 202         public:
 203                 Heap()
 204                 {
 205                         top = NULL;
 206                         num_nodes = 0;
 207                 }
 208
 209                 bool empty()
 210                 {
 211                         return num_nodes == 0;
 212                 }
 213
 214                 T *get_top()
 215                 {
 216                         if (!top)
 217                                 return NULL;
 218
 219                         return node_to_type(top);
 220                 }
 221
 222                 void add(T *t)
 223                 {
 224                         Node *n = &t->heap_node;
 225                         assert(!n->is_on_heap());
 226
 227                         Node *ret = get_by_index(num_nodes, &n->parent, &n->parentlink);
 228
 229                         assert(ret == NULL);
 230                         assert(num_nodes == 0 || n->parent);
 231
 232                         num_nodes++;
 233                         *n->parentlink = n;
 234                         n->left = n->right = NULL;
 235
 236                         up(n);
 237                 }
 238
 239                 void del(T *t)
 240                 {
 241                         Node *n = &t->heap_node;
 242                         assert(*n->parentlink == n);
 243
 244                         if (--num_nodes == 0) {
 245                                 top = NULL;
 246                                 n->parentlink = NULL;
 247                                 return;
 248                         }
 249
 250                         Node *parent, **link;
 251                         Node *last = get_by_index(num_nodes, &parent, &link);
 252
 253                         assert(last->parent);
 254                         assert(last->is_on_heap());
 255                         assert(!last->left);
 256                         assert(!last->right);
 257
 258                         *link = NULL;
 259                         *n->parentlink = last;
 260                         *last = *n;
 261                         n->parentlink = NULL;
 262
 263                         if (last->left) {
 264                                 last->left->parent = last;
 265                                 last->left->parentlink = &last->left;
 266                         }
 267
 268                         if (last->right) {
 269                                 last->right->parent = last;
 270                                 last->right->parentlink = &last->right;
 271                         }
 272
 273                         down(last);
 274                 }
 275
 276                 void requeue(T *t) {
 277                         Node *n = &t->heap_node;
 278                         assert(*n->parentlink == n);
 279
 280                         if (!n->parent || *node_to_type(n->parent) < *t)
 281                                 down(n);
 282                         else
 283                                 up(n);
 284                 }
 285         };
 286 }
 287
 288 #endif