c_cpp/lib/klib/cpp/kavl.hpp

   1 #ifndef KAVL_HPP
   2 #define KAVL_HPP
   3
   4 #include <functional>
   5
   6 namespace klib {
   7
   8 template<class T, typename Less = std::less<T> >
   9 class Avl {
  10         static const int MAX_DEPTH = 64;
  11         struct Node {
  12                 T data;
  13                 signed char balance;
  14                 unsigned size;
  15                 Node *p[2];
  16         };
  17         Node *root;
  18         inline int cmp_func(const T &x, const T &y) {
  19                 return Less()(y, x) - Less()(x, y);
  20         }
  21         inline unsigned child_size(Node *p, int dir) {
  22                 return p->p[dir]? p->p[dir]->size : 0;
  23         };
  24         // one rotation: (a,(b,c)q)p => ((a,b)p,c)q
  25         inline Node *rotate1(Node *p, int dir) { // dir=0 to left; dir=1 to right
  26                 int opp = 1 - dir; // opposite direction
  27                 Node *q = p->p[opp];
  28                 unsigned size_p = p->size;
  29                 p->size -= q->size - child_size(q, dir);
  30                 q->size = size_p;
  31                 p->p[opp] = q->p[dir];
  32                 q->p[dir] = p;
  33                 return q;
  34         };
  35         // two consecutive rotations: (a,((b,c)r,d)q)p => ((a,b)p,(c,d)q)r
  36         inline Node *rotate2(Node *p, int dir) {
  37                 int b1, opp = 1 - dir;
  38                 Node *q = p->p[opp], *r = q->p[dir];
  39                 unsigned size_x_dir = child_size(r, dir);
  40                 r->size = p->size;
  41                 p->size -= q->size - size_x_dir;
  42                 q->size -= size_x_dir + 1;
  43                 p->p[opp] = r->p[dir];
  44                 r->p[dir] = p;
  45                 q->p[dir] = r->p[opp];
  46                 r->p[opp] = q;
  47                 b1 = dir == 0? +1 : -1;
  48                 if (r->balance == b1) q->balance = 0, p->balance = -b1;
  49                 else if (r->balance == 0) q->balance = p->balance = 0;
  50                 else q->balance = b1, p->balance = 0;
  51                 r->balance = 0;
  52                 return r;
  53         };
  54         void destroy(Node *r) {
  55                 Node *p, *q;
  56                 for (p = r; p; p = q) {
  57                         if (p->p[0] == 0) {
  58                                 q = p->p[1];
  59                                 delete p;
  60                         } else {
  61                                 q = p->p[0];
  62                                 p->p[0] = q->p[1];
  63                                 q->p[1] = p;
  64                         }
  65                 }
  66         };
  67 public:
  68         Avl() : root(NULL) {};
  69         ~Avl() { destroy(root); };
  70         unsigned size() const { return root? root->size : 0; }
  71         T *find(const T &data, unsigned *cnt_ = NULL) {
  72                 Node *p = root;
  73                 unsigned cnt = 0;
  74                 while (p != 0) {
  75                         int cmp = cmp_func(data, p->data);
  76                         if (cmp >= 0) cnt += child_size(p, 0) + 1;
  77                         if (cmp < 0) p = p->p[0];
  78                         else if (cmp > 0) p = p->p[1];
  79                         else break;
  80                 }
  81                 if (cnt_) *cnt_ = cnt;
  82                 return p? &p->data : NULL;
  83         };
  84         T *insert(const T &data, bool *is_new = NULL, unsigned *cnt_ = NULL) {
  85                 unsigned char stack[MAX_DEPTH];
  86                 Node *path[MAX_DEPTH];
  87                 Node *bp, *bq;
  88                 Node *x, *p, *q, *r = 0; // _r_ is potentially the new root
  89                 int i, which = 0, top, b1, path_len;
  90                 unsigned cnt = 0;
  91                 bp = root, bq = 0;
  92                 if (is_new) *is_new = false;
  93                 // find the insertion location
  94                 for (p = bp, q = bq, top = path_len = 0; p; q = p, p = p->p[which]) {
  95                         int cmp = cmp_func(data, p->data);
  96                         if (cmp >= 0) cnt += child_size(p, 0) + 1;
  97                         if (cmp == 0) {
  98                                 if (cnt_) *cnt_ = cnt;
  99                                 return &p->data;
 100                         }
 101                         if (p->balance != 0)
 102                                 bq = q, bp = p, top = 0;
 103                         stack[top++] = which = (cmp > 0);
 104                         path[path_len++] = p;
 105                 }
 106                 if (cnt_) *cnt_ = cnt;
 107                 x = new Node;
 108                 x->data = data, x->balance = 0, x->size = 1, x->p[0] = x->p[1] = 0;
 109                 if (is_new) *is_new = true;
 110                 if (q == 0) root = x;
 111                 else q->p[which] = x;
 112                 if (bp == 0) return &x->data;
 113                 for (i = 0; i < path_len; ++i) ++path[i]->size;
 114                 for (p = bp, top = 0; p != x; p = p->p[stack[top]], ++top) /* update balance factors */
 115                         if (stack[top] == 0) --p->balance;
 116                         else ++p->balance;
 117                 if (bp->balance > -2 && bp->balance < 2) return &x->data; /* no re-balance needed */
 118                 // re-balance
 119                 which = (bp->balance < 0);
 120                 b1 = which == 0? +1 : -1;
 121                 q = bp->p[1 - which];
 122                 if (q->balance == b1) {
 123                         r = rotate1(bp, which);
 124                         q->balance = bp->balance = 0;
 125                 } else r = rotate2(bp, which);
 126                 if (bq == 0) root = r;
 127                 else bq->p[bp != bq->p[0]] = r;
 128                 return &x->data;
 129         };
 130         bool erase(const T &data) {
 131                 Node *p, *path[MAX_DEPTH], fake;
 132                 unsigned char dir[MAX_DEPTH];
 133                 int i, d = 0, cmp;
 134                 fake.p[0] = root, fake.p[1] = 0;
 135                 for (cmp = -1, p = &fake; cmp; cmp = cmp_func(data, p->data)) {
 136                         int which = (cmp > 0);
 137                         dir[d] = which;
 138                         path[d++] = p;
 139                         p = p->p[which];
 140                         if (p == 0) return false;
 141                 }
 142                 for (i = 1; i < d; ++i) --path[i]->size;
 143                 if (p->p[1] == 0) { // ((1,.)2,3)4 => (1,3)4; p=2
 144                         path[d-1]->p[dir[d-1]] = p->p[0];
 145                 } else {
 146                         Node *q = p->p[1];
 147                         if (q->p[0] == 0) { // ((1,2)3,4)5 => ((1)2,4)5; p=3
 148                                 q->p[0] = p->p[0];
 149                                 q->balance = p->balance;
 150                                 path[d-1]->p[dir[d-1]] = q;
 151                                 path[d] = q, dir[d++] = 1;
 152                                 q->size = p->size - 1;
 153                         } else { // ((1,((.,2)3,4)5)6,7)8 => ((1,(2,4)5)3,7)8; p=6
 154                                 Node *r;
 155                                 int e = d++; // backup _d_
 156                                 for (;;) {
 157                                         dir[d] = 0;
 158                                         path[d++] = q;
 159                                         r = q->p[0];
 160                                         if (r->p[0] == 0) break;
 161                                         q = r;
 162                                 }
 163                                 r->p[0] = p->p[0];
 164                                 q->p[0] = r->p[1];
 165                                 r->p[1] = p->p[1];
 166                                 r->balance = p->balance;
 167                                 path[e-1]->p[dir[e-1]] = r;
 168                                 path[e] = r, dir[e] = 1;
 169                                 for (i = e + 1; i < d; ++i) --path[i]->size;
 170                                 r->size = p->size - 1;
 171                         }
 172                 }
 173                 while (--d > 0) {
 174                         Node *q = path[d];
 175                         int which, other, b1 = 1, b2 = 2;
 176                         which = dir[d], other = 1 - which;
 177                         if (which) b1 = -b1, b2 = -b2;
 178                         q->balance += b1;
 179                         if (q->balance == b1) break;
 180                         else if (q->balance == b2) {
 181                                 Node *r = q->p[other];
 182                                 if (r->balance == -b1) {
 183                                         path[d-1]->p[dir[d-1]] = rotate2(q, which);
 184                                 } else {
 185                                         path[d-1]->p[dir[d-1]] = rotate1(q, which);
 186                                         if (r->balance == 0) {
 187                                                 r->balance = -b1;
 188                                                 q->balance = b1;
 189                                                 break;
 190                                         } else r->balance = q->balance = 0;
 191                                 }
 192                         }
 193                 }
 194                 root = fake.p[0];
 195                 delete p;
 196                 return true;
 197         };
 198 };
 199
 200 } // end of namespace klib
 201
 202 #endif