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