avl.c (6857B)
1 #include <u.h> 2 #include <libc.h> 3 #include <bio.h> 4 #include <avl.h> 5 6 /* 7 * In-memory database stored as self-balancing AVL tree. 8 * See Lewis & Denenberg, Data Structures and Their Algorithms. 9 */ 10 11 static void 12 singleleft(Avl **tp, Avl *p) 13 { 14 int l, r2; 15 Avl *a, *c; 16 17 a = *tp; 18 c = a->n[1]; 19 20 r2 = c->bal; 21 l = (r2 > 0? r2: 0)+1 - a->bal; 22 23 if((a->n[1] = c->n[0]) != nil) 24 a->n[1]->p = a; 25 26 if((c->n[0] = a) != nil) 27 c->n[0]->p = c; 28 29 if((*tp = c) != nil) 30 (*tp)->p = p; 31 32 a->bal = -l; 33 c->bal = r2 - ((l > 0? l: 0)+1); 34 35 } 36 37 static void 38 singleright(Avl **tp, Avl *p) 39 { 40 int l2, r; 41 Avl *a, *c; 42 43 a = *tp; 44 c = a->n[0]; 45 l2 = - c->bal; 46 r = a->bal + ((l2 > 0? l2: 0)+1); 47 48 if((a->n[0] = c->n[1]) != nil) 49 a->n[0]->p = a; 50 51 if((c->n[1] = a) != nil) 52 c->n[1]->p = c; 53 54 if((*tp = c) != nil) 55 (*tp)->p = p; 56 57 a->bal = r; 58 c->bal = ((r > 0? r: 0)+1) - l2; 59 } 60 61 static void 62 doublerightleft(Avl **tp, Avl *p) 63 { 64 singleright(&(*tp)->n[1], *tp); 65 singleleft(tp, p); 66 } 67 68 static void 69 doubleleftright(Avl **tp, Avl *p) 70 { 71 singleleft(&(*tp)->n[0], *tp); 72 singleright(tp, p); 73 } 74 75 static void 76 balance(Avl **tp, Avl *p) 77 { 78 switch((*tp)->bal){ 79 case -2: 80 if((*tp)->n[0]->bal <= 0) 81 singleright(tp, p); 82 else if((*tp)->n[0]->bal == 1) 83 doubleleftright(tp, p); 84 else 85 assert(0); 86 break; 87 88 case 2: 89 if((*tp)->n[1]->bal >= 0) 90 singleleft(tp, p); 91 else if((*tp)->n[1]->bal == -1) 92 doublerightleft(tp, p); 93 else 94 assert(0); 95 break; 96 } 97 } 98 99 static int 100 canoncmp(int cmp) 101 { 102 if(cmp < 0) 103 return -1; 104 else if(cmp > 0) 105 return 1; 106 return 0; 107 } 108 109 static int 110 _insertavl(Avl **tp, Avl *p, Avl *r, int (*cmp)(Avl*,Avl*), Avl **rfree) 111 { 112 int i, ob; 113 114 if(*tp == nil){ 115 r->bal = 0; 116 r->n[0] = nil; 117 r->n[1] = nil; 118 r->p = p; 119 *tp = r; 120 return 1; 121 } 122 ob = (*tp)->bal; 123 if((i = canoncmp(cmp(r, *tp))) != 0){ 124 (*tp)->bal += i * _insertavl(&(*tp)->n[(i+1)/2], *tp, r, cmp, 125 rfree); 126 balance(tp, p); 127 return ob == 0 && (*tp)->bal != 0; 128 } 129 130 /* install new entry */ 131 *rfree = *tp; /* save old node for freeing */ 132 *tp = r; /* insert new node */ 133 **tp = **rfree; /* copy old node's Avl contents */ 134 if(r->n[0]) /* fix node's children's parent pointers */ 135 r->n[0]->p = r; 136 if(r->n[1]) 137 r->n[1]->p = r; 138 139 return 0; 140 } 141 142 static int 143 successor(Avl **tp, Avl *p, Avl **r) 144 { 145 int ob; 146 147 if((*tp)->n[0] == nil){ 148 *r = *tp; 149 *tp = (*r)->n[1]; 150 if(*tp) 151 (*tp)->p = p; 152 return -1; 153 } 154 ob = (*tp)->bal; 155 (*tp)->bal -= successor(&(*tp)->n[0], *tp, r); 156 balance(tp, p); 157 return -(ob != 0 && (*tp)->bal == 0); 158 } 159 160 static int 161 _deleteavl(Avl **tp, Avl *p, Avl *rx, int(*cmp)(Avl*,Avl*), Avl **del, 162 void (*predel)(Avl*, void*), void *arg) 163 { 164 int i, ob; 165 Avl *r, *or; 166 167 if(*tp == nil) 168 return 0; 169 170 ob = (*tp)->bal; 171 if((i=canoncmp(cmp(rx, *tp))) != 0){ 172 (*tp)->bal += i * _deleteavl(&(*tp)->n[(i+1)/2], *tp, rx, cmp, 173 del, predel, arg); 174 balance(tp, p); 175 return -(ob != 0 && (*tp)->bal == 0); 176 } 177 178 if(predel) 179 (*predel)(*tp, arg); 180 181 or = *tp; 182 if(or->n[i=0] == nil || or->n[i=1] == nil){ 183 *tp = or->n[1-i]; 184 if(*tp) 185 (*tp)->p = p; 186 *del = or; 187 return -1; 188 } 189 190 /* deleting node with two kids, find successor */ 191 or->bal += successor(&or->n[1], or, &r); 192 r->bal = or->bal; 193 r->n[0] = or->n[0]; 194 r->n[1] = or->n[1]; 195 *tp = r; 196 (*tp)->p = p; 197 /* node has changed; fix children's parent pointers */ 198 if(r->n[0]) 199 r->n[0]->p = r; 200 if(r->n[1]) 201 r->n[1]->p = r; 202 *del = or; 203 balance(tp, p); 204 return -(ob != 0 && (*tp)->bal == 0); 205 } 206 207 /* 208 static void 209 checkparents(Avl *a, Avl *p) 210 { 211 if(a == nil) 212 return; 213 if(a->p != p) 214 print("bad parent\n"); 215 checkparents(a->n[0], a); 216 checkparents(a->n[1], a); 217 } 218 */ 219 220 struct Avltree 221 { 222 Avl *root; 223 int (*cmp)(Avl*, Avl*); 224 Avlwalk *walks; 225 }; 226 struct Avlwalk 227 { 228 int started; 229 int moved; 230 Avlwalk *next; 231 Avltree *tree; 232 Avl *node; 233 }; 234 235 Avltree* 236 mkavltree(int (*cmp)(Avl*, Avl*)) 237 { 238 Avltree *t; 239 240 t = malloc(sizeof *t); 241 if(t == nil) 242 return nil; 243 memset(t, 0, sizeof *t); 244 t->cmp = cmp; 245 return t; 246 } 247 248 void 249 insertavl(Avltree *t, Avl *new, Avl **oldp) 250 { 251 *oldp = nil; 252 _insertavl(&t->root, nil, new, t->cmp, oldp); 253 } 254 255 static Avl* 256 findpredecessor(Avl *a) 257 { 258 if(a == nil) 259 return nil; 260 261 if(a->n[0] != nil){ 262 /* predecessor is rightmost descendant of left child */ 263 for(a = a->n[0]; a->n[1]; a = a->n[1]) 264 ; 265 return a; 266 }else{ 267 /* we're at a leaf, successor is a parent we enter from the right */ 268 while(a->p && a->p->n[0] == a) 269 a = a->p; 270 return a->p; 271 } 272 } 273 274 static Avl* 275 findsuccessor(Avl *a) 276 { 277 if(a == nil) 278 return nil; 279 280 if(a->n[1] != nil){ 281 /* successor is leftmost descendant of right child */ 282 for(a = a->n[1]; a->n[0]; a = a->n[0]) 283 ; 284 return a; 285 }else{ 286 /* we're at a leaf, successor is a parent we enter from the left going up */ 287 while(a->p && a->p->n[1] == a) 288 a = a->p; 289 return a->p; 290 } 291 } 292 293 static Avl* 294 _lookupavl(Avl *t, Avl *r, int (*cmp)(Avl*,Avl*), int neighbor) 295 { 296 int i; 297 Avl *p; 298 299 p = nil; 300 if(t == nil) 301 return nil; 302 do{ 303 assert(t->p == p); 304 if((i = canoncmp(cmp(r, t))) == 0) 305 return t; 306 p = t; 307 t = t->n[(i+1)/2]; 308 }while(t); 309 if(neighbor == 0) 310 return nil; 311 if(neighbor < 0) 312 return i > 0 ? p : findpredecessor(p); 313 return i < 0 ? p : findsuccessor(p); 314 } 315 316 Avl* 317 searchavl(Avltree *t, Avl *key, int neighbor) 318 { 319 return _lookupavl(t->root, key, t->cmp, neighbor); 320 } 321 322 Avl* 323 lookupavl(Avltree *t, Avl *key) 324 { 325 return _lookupavl(t->root, key, t->cmp, 0); 326 } 327 328 static void 329 walkdel(Avl *a, void *v) 330 { 331 Avl *p; 332 Avlwalk *w; 333 Avltree *t; 334 335 if(a == nil) 336 return; 337 338 p = findpredecessor(a); 339 t = v; 340 for(w = t->walks; w; w = w->next){ 341 if(w->node == a){ 342 /* back pointer to predecessor; not perfect but adequate */ 343 w->moved = 1; 344 w->node = p; 345 if(p == nil) 346 w->started = 0; 347 } 348 } 349 } 350 351 void 352 deleteavl(Avltree *t, Avl *key, Avl **oldp) 353 { 354 *oldp = nil; 355 _deleteavl(&t->root, nil, key, t->cmp, oldp, walkdel, t); 356 } 357 358 Avlwalk* 359 avlwalk(Avltree *t) 360 { 361 Avlwalk *w; 362 363 w = malloc(sizeof *w); 364 if(w == nil) 365 return nil; 366 memset(w, 0, sizeof *w); 367 w->tree = t; 368 w->next = t->walks; 369 t->walks = w; 370 return w; 371 } 372 373 Avl* 374 avlnext(Avlwalk *w) 375 { 376 Avl *a; 377 378 if(w->started==0){ 379 for(a = w->tree->root; a && a->n[0]; a = a->n[0]) 380 ; 381 w->node = a; 382 w->started = 1; 383 }else{ 384 a = findsuccessor(w->node); 385 if(a == w->node) 386 abort(); 387 w->node = a; 388 } 389 return w->node; 390 } 391 392 Avl* 393 avlprev(Avlwalk *w) 394 { 395 Avl *a; 396 397 if(w->started == 0){ 398 for(a = w->tree->root; a && a->n[1]; a = a->n[1]) 399 ; 400 w->node = a; 401 w->started = 1; 402 }else if(w->moved){ 403 w->moved = 0; 404 return w->node; 405 }else{ 406 a = findpredecessor(w->node); 407 if(a == w->node) 408 abort(); 409 w->node = a; 410 } 411 return w->node; 412 } 413 414 void 415 endwalk(Avlwalk *w) 416 { 417 Avltree *t; 418 Avlwalk **l; 419 420 t = w->tree; 421 for(l = &t->walks; *l; l = &(*l)->next){ 422 if(*l == w){ 423 *l = w->next; 424 break; 425 } 426 } 427 free(w); 428 } 429 430 /* 431 static void 432 walkavl(Avl *t, void (*f)(Avl*, void*), void *v) 433 { 434 if(t == nil) 435 return; 436 walkavl(t->n[0], f, v); 437 f(t, v); 438 walkavl(t->n[1], f, v); 439 } 440 */