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