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 // Permission is hereby granted, free of charge, to any person obtaining a copy of
   6 // this software and associated documentation files (the "Software"), to deal with
   7 // the Software without restriction, including without limitation the rights to
   8 // use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
   9 // of the Software, and to permit persons to whom the Software is furnished to do
  10 // so, subject to the following condition:
  11 //
  12 // The above copyright notice and this permission notice shall be
  13 // included in all copies or substantial portions of the Software.
  14 //
  15 // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  16 // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
  17 // FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
  18 // CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  19 // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  20 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS WITH THE
  21 // SOFTWARE.
  22
  23 #ifndef _UTIL_HEAP_H
  24 #define _UTIL_HEAP_H
  25
  26 #include <assert.h>
  27 #include <stddef.h>
  28 #include <stdint.h>
  29 #include <lowlevel/bitops.h>
  30
  31 namespace Util {
  32         // T must support operator < for key ordering, and
  33         // have a Heap<T>::Node member called heap_node.
  34
  35         template <typename T>
  36         class Heap {
  37         public:
  38                 struct Node {
  39                         Node *parent, *left, *right;
  40
  41                         // A pointer to the parent's left or right pointer
  42                         // (whichever side this node is on).  This is also
  43                         // used for sanity checks (it is NULL for nodes not
  44                         // on a queue).
  45
  46                         Node **parentlink;
  47
  48                         Node()
  49                         {
  50                                 parent = left = right = NULL;
  51                                 parentlink = NULL;
  52                         }
  53
  54                         bool is_on_heap()
  55                         {
  56                                 return parentlink != NULL;
  57                         }
  58                 };
  59
  60         private:
  61                 Node *top;
  62                 size_t num_nodes;
  63
  64                 T *node_to_type(Node *n)
  65                 {
  66                         return reinterpret_cast<T *>(reinterpret_cast<uintptr_t>(n) -
  67                                                      offsetof(T, heap_node));
  68                 }
  69
  70                 // Swap n with its parent
  71                 void swap(Node *n)
  72                 {
  73                         Node *p = n->parent;
  74                         *p->parentlink = n;
  75
  76                         n->parent = p->parent;
  77                         p->parent = n;
  78
  79                         Node **plink = p->parentlink;
  80
  81                         if (n->parentlink == &p->left) {
  82                                 assert(p->left == n);
  83
  84                                 Node *tmp = p->right;
  85                                 p->right = n->right;
  86                                 n->right = tmp;
  87
  88                                 p->left = n->left;
  89                                 n->left = p;
  90
  91                                 p->parentlink = &n->left;
  92
  93                                 if (n->right) {
  94                                         n->right->parent = n;
  95                                         n->right->parentlink = &n->right;
  96                                 }
  97                         } else {
  98                                 assert(n->parentlink == &p->right);
  99                                 assert(p->right == n);
 100
 101                                 Node *tmp = p->left;
 102                                 p->left = n->left;
 103                                 n->left = tmp;
 104
 105                                 p->right = n->right;
 106                                 n->right = p;
 107
 108                                 p->parentlink = &n->right;
 109
 110                                 if (n->left) {
 111                                         n->left->parent = n;
 112                                         n->left->parentlink = &n->left;
 113                                 }
 114                         }
 115
 116                         if (p->left) {
 117                                 p->left->parent = p;
 118                                 p->left->parentlink = &p->left;
 119
 120                                 if (p->right) {
 121                                         p->right->parent = p;
 122                                         p->right->parentlink = &p->right;
 123                                 }
 124                         } else {
 125                                 assert(!p->right);
 126                         }
 127
 128                         n->parentlink = plink;
 129                 }
 130
 131                 // Restore heap condition by moving n towards the root,
 132                 // if required.
 133
 134                 void up(Node *n)
 135                 {
 136                         T *n_type = node_to_type(n);
 137
 138                         assert(!n->left || !(*node_to_type(n->left) < *n_type));
 139                         assert(!n->right || !(*node_to_type(n->right) < *n_type));
 140
 141                         while (n->parent && *n_type < *node_to_type(n->parent)) {
 142                                 assert(!n->left || !(*node_to_type(n->left) < *n_type));
 143                                 assert(!n->right || !(*node_to_type(n->right) < *n_type));
 144
 145                                 swap(n);
 146                         }
 147
 148                         assert(!n->left || !(*node_to_type(n->left) < *n_type));
 149                         assert(!n->right || !(*node_to_type(n->right) < *n_type));
 150                 }
 151
 152                 // Restore heap condition by moving n towards the leaves,
 153                 // if required.
 154
 155                 void down(Node *n)
 156                 {
 157                         T *n_type = node_to_type(n);
 158
 159                         assert(!n->parent || !(*n_type < *node_to_type(n->parent)));
 160
 161                         // If n has a right, it'll have a left.
 162                         while (n->left) {
 163                                 assert(!n->parent || !(*n_type < *node_to_type(n->parent)));
 164                                 Node *child = n->left;
 165
 166                                 if (n->right &&
 167                                     *node_to_type(n->right) < *node_to_type(n->left))
 168                                         child = n->right;
 169
 170                                 if (*node_to_type(child) < *n_type)
 171                                         swap(child);
 172                                 else
 173                                         break;
 174                         }
 175
 176                         assert(!n->parent || !(*n_type < *node_to_type(n->parent)));
 177                 }
 178
 179                 Node *get_by_index(size_t num, Node **parent, Node ***link)
 180                 {
 181                         assert(num_nodes == 0 || top);
 182                         assert(num <= num_nodes);
 183
 184                         // It's easier to use 1-based indices, so that the high set bit
 185                         // perfectly indicates which level the node is on.
 186
 187                         num += 1;
 188
 189                         Node *n = NULL;
 190                         size_t bit = 1 << ll_get_order_round_down(num);
 191                         *parent = NULL;
 192                         Node **nextlink = &top;
 193
 194                         do {
 195                                 *link = nextlink;
 196                                 *parent = n;
 197                                 n = *nextlink;
 198
 199                                 bit /= 2;
 200
 201                                 if (num & bit)
 202                                         nextlink = &n->right;
 203                                 else
 204                                         nextlink = &n->left;
 205                         } while (bit);
 206
 207                         return n;
 208                 }
 209
 210         public:
 211                 Heap()
 212                 {
 213                         top = NULL;
 214                         num_nodes = 0;
 215                 }
 216
 217                 bool empty()
 218                 {
 219                         return num_nodes == 0;
 220                 }
 221
 222                 T *get_top()
 223                 {
 224                         if (!top)
 225                                 return NULL;
 226
 227                         return node_to_type(top);
 228                 }
 229
 230                 void add(T *t)
 231                 {
 232                         Node *n = &t->heap_node;
 233                         assert(!n->is_on_heap());
 234
 235                         Node *ret = get_by_index(num_nodes, &n->parent, &n->parentlink);
 236
 237                         assert(ret == NULL);
 238                         assert(num_nodes == 0 || n->parent);
 239
 240                         num_nodes++;
 241                         *n->parentlink = n;
 242                         n->left = n->right = NULL;
 243
 244                         up(n);
 245                 }
 246
 247                 void del(T *t)
 248                 {
 249                         Node *n = &t->heap_node;
 250                         assert(*n->parentlink == n);
 251
 252                         if (--num_nodes == 0) {
 253                                 top = NULL;
 254                                 n->parentlink = NULL;
 255                                 return;
 256                         }
 257
 258                         Node *parent, **link;
 259                         Node *last = get_by_index(num_nodes, &parent, &link);
 260
 261                         assert(last->parent);
 262                         assert(last->is_on_heap());
 263                         assert(!last->left);
 264                         assert(!last->right);
 265
 266                         *link = NULL;
 267                         *n->parentlink = last;
 268                         *last = *n;
 269                         n->parentlink = NULL;
 270
 271                         if (last->left) {
 272                                 last->left->parent = last;
 273                                 last->left->parentlink = &last->left;
 274                         }
 275
 276                         if (last->right) {
 277                                 last->right->parent = last;
 278                                 last->right->parentlink = &last->right;
 279                         }
 280
 281                         down(last);
 282                 }
 283
 284                 void requeue(T *t) {
 285                         Node *n = &t->heap_node;
 286                         assert(*n->parentlink == n);
 287
 288                         if (!n->parent || *node_to_type(n->parent) < *t)
 289                                 down(n);
 290                         else
 291                                 up(n);
 292                 }
 293         };
 294 }
 295
 296 #endif