sys/src/libavl/avl.c

   1 #include <u.h>
   2 #include <libc.h>
   3 #include <avl.h>
   4
   5 /* See Knuth Volume 3, 6.2.3 */
   6
   7 Avl *avllookup(Avltree*, Avl*);
   8 Avl *avldelete(Avltree*, Avl*);
   9 Avl *avlinsert(Avltree*, Avl*);
  10
  11 Avltree*
  12 avlcreate(int (*cmp)(Avl*, Avl*))
  13 {
  14         Avltree *t;
  15
  16         t = malloc(sizeof(*t));
  17         if(t == nil)
  18                 return nil;
  19
  20         t->cmp = cmp;
  21         t->root = nil;
  22         return t;
  23 }
  24
  25 Avl*
  26 avllookup(Avltree *t, Avl *k)
  27 {
  28         Avl *h;
  29         int c;
  30
  31         h = t->root;
  32         while(h != nil){
  33                 c = (t->cmp)(k, h);
  34                 if(c < 0){
  35                         h = h->c[0];
  36                         continue;
  37                 }
  38                 if(c > 0){
  39                         h = h->c[1];
  40                         continue;
  41                 }
  42                 return h;
  43         }
  44         return nil;
  45 }
  46
  47 static int insert(int (*)(Avl*, Avl*), Avl*, Avl**, Avl*, Avl**);
  48
  49 Avl*
  50 avlinsert(Avltree *t, Avl *k)
  51 {
  52         Avl *old;
  53
  54         old = nil;
  55         insert(t->cmp, nil, &t->root, k, &old);
  56         return old;
  57 }
  58
  59 static int insertfix(int, Avl**);
  60
  61 static int
  62 insert(int (*cmp)(Avl*, Avl*), Avl *p, Avl **qp, Avl *k, Avl **oldp)
  63 {
  64         Avl *q;
  65         int fix, c;
  66
  67         q = *qp;
  68         if(q == nil) {
  69                 k->c[0] = nil;
  70                 k->c[1] = nil;
  71                 k->balance = 0;
  72                 k->p = p;
  73                 *qp = k;
  74                 return 1;
  75         }
  76
  77         c = cmp(k, q);
  78         c = c > 0 ? 1 : c < 0 ? -1: 0;
  79         if(c == 0) {
  80                 *oldp = q;
  81                 *k = *q;
  82                 if(q->c[0] != nil)
  83                         q->c[0]->p = k;
  84                 if(q->c[1] != nil)
  85                         q->c[1]->p = k;
  86                 *qp = k;
  87                 return 0;
  88         }
  89         fix = insert(cmp, q, q->c + (c+1)/2, k, oldp);
  90         if(fix)
  91                 return insertfix(c, qp);
  92         return 0;
  93 }
  94
  95 static Avl *singlerot(int, Avl*);
  96 static Avl *doublerot(int, Avl*);
  97
  98 static int
  99 insertfix(int c, Avl **t)
 100 {
 101         Avl *s;
 102
 103         s = *t;
 104         if(s->balance == 0) {
 105                 s->balance = c;
 106                 return 1;
 107         }
 108         if(s->balance == -c) {
 109                 s->balance = 0;
 110                 return 0;
 111         }
 112         if(s->c[(c+1)/2]->balance == c)
 113                 s = singlerot(c, s);
 114         else
 115                 s = doublerot(c, s);
 116         *t = s;
 117         return 0;
 118 }
 119
 120 static int delete(int (*cmp)(Avl*, Avl*), Avl**, Avl*, Avl**);
 121 static int deletemin(Avl**, Avl**);
 122 static int deletefix(int, Avl**);
 123
 124 Avl*
 125 avldelete(Avltree *t, Avl *k)
 126 {
 127         Avl *old;
 128
 129         if(t->root == nil)
 130                 return nil;
 131         old = nil;
 132         delete(t->cmp, &t->root, k, &old);
 133         return old;
 134 }
 135
 136 static int
 137 delete(int (*cmp)(Avl*, Avl*), Avl **qp, Avl *k, Avl **oldp)
 138 {
 139         Avl *q, *e;
 140         int c, fix;
 141
 142         q = *qp;
 143         if(q == nil)
 144                 return 0;
 145
 146         c = cmp(k, q);
 147         c = c > 0 ? 1 : c < 0 ? -1: 0;
 148         if(c == 0) {
 149                 *oldp = q;
 150                 if(q->c[1] == nil) {
 151                         *qp = q->c[0];
 152                         if(*qp != nil)
 153                                 (*qp)->p = q->p;
 154                         return 1;
 155                 }
 156                 fix = deletemin(q->c+1, &e);
 157                 *e = *q;
 158                 if(q->c[0] != nil)
 159                         q->c[0]->p = e;
 160                 if(q->c[1] != nil)
 161                         q->c[1]->p = e;
 162                 *qp = e;
 163                 if(fix)
 164                         return deletefix(-1, qp);
 165                 return 0;
 166         }
 167         fix = delete(cmp, q->c + (c+1)/2, k, oldp);
 168         if(fix)
 169                 return deletefix(-c, qp);
 170         return 0;
 171 }
 172
 173 static int
 174 deletemin(Avl **qp, Avl **oldp)
 175 {
 176         Avl *q;
 177         int fix;
 178
 179         q = *qp;
 180         if(q->c[0] == nil) {
 181                 *oldp = q;
 182                 *qp = q->c[1];
 183                 if(*qp != nil)
 184                         (*qp)->p = q->p;
 185                 return 1;
 186         }
 187         fix = deletemin(q->c, oldp);
 188         if(fix)
 189                 return deletefix(1, qp);
 190         return 0;
 191 }
 192
 193 static Avl *rotate(int, Avl*);
 194
 195 static int
 196 deletefix(int c, Avl **t)
 197 {
 198         Avl *s;
 199         int a;
 200
 201         s = *t;
 202         if(s->balance == 0) {
 203                 s->balance = c;
 204                 return 0;
 205         }
 206         if(s->balance == -c) {
 207                 s->balance = 0;
 208                 return 1;
 209         }
 210         a = (c+1)/2;
 211         if(s->c[a]->balance == 0) {
 212                 s = rotate(c, s);
 213                 s->balance = -c;
 214                 *t = s;
 215                 return 0;
 216         }
 217         if(s->c[a]->balance == c)
 218                 s = singlerot(c, s);
 219         else
 220                 s = doublerot(c, s);
 221         *t = s;
 222         return 1;
 223 }
 224
 225 static Avl*
 226 singlerot(int c, Avl *s)
 227 {
 228         s->balance = 0;
 229         s = rotate(c, s);
 230         s->balance = 0;
 231         return s;
 232 }
 233
 234 static Avl*
 235 doublerot(int c, Avl *s)
 236 {
 237         Avl *r, *p;
 238         int a;
 239
 240         a = (c+1)/2;
 241         r = s->c[a];
 242         s->c[a] = rotate(-c, s->c[a]);
 243         p = rotate(c, s);
 244         assert(r->p == p);
 245         assert(s->p == p);
 246
 247         if(p->balance == c) {
 248                 s->balance = -c;
 249                 r->balance = 0;
 250         } else if(p->balance == -c) {
 251                 s->balance = 0;
 252                 r->balance = c;
 253         } else
 254                 s->balance = r->balance = 0;
 255         p->balance = 0;
 256         return p;
 257 }
 258
 259 static Avl*
 260 rotate(int c, Avl *s)
 261 {
 262         Avl *r, *n;
 263         int a;
 264
 265         a = (c+1)/2;
 266         r = s->c[a];
 267         s->c[a] = n = r->c[a^1];
 268         if(n != nil)
 269                 n->p = s;
 270         r->c[a^1] = s;
 271         r->p = s->p;
 272         s->p = r;
 273         return r;
 274 }
 275
 276 static Avl *walk1(int, Avl*);
 277
 278 Avl*
 279 avlprev(Avl *q)
 280 {
 281         return walk1(0, q);
 282 }
 283
 284 Avl*
 285 avlnext(Avl *q)
 286 {
 287         return walk1(1, q);
 288 }
 289
 290 static Avl*
 291 walk1(int a, Avl *q)
 292 {
 293         Avl *p;
 294
 295         if(q == nil)
 296                 return nil;
 297
 298         if(q->c[a] != nil){
 299                 for(q = q->c[a]; q->c[a^1] != nil; q = q->c[a^1])
 300                         ;
 301                 return q;
 302         }
 303         for(p = q->p; p != nil && p->c[a] == q; p = p->p)
 304                 q = p;
 305         return p;
 306 }