| 1 | #include "avl.h"
 | 
|---|
| 2 | #include "avl-private.h"
 | 
|---|
| 3 | 
 | 
|---|
| 4 | // AVL tree specific (internal) operations:
 | 
|---|
| 5 | // rotateLeft, rotateRight, fix
 | 
|---|
| 6 | //
 | 
|---|
| 7 | // AVL tree enhanced height operation
 | 
|---|
| 8 | //
 | 
|---|
| 9 | // calcBalance is a simple computation of height(R) - height(L)
 | 
|---|
| 10 | 
 | 
|---|
| 11 | // an AVL tree's height is easy to compute
 | 
|---|
| 12 | // just follow path with the larger balance
 | 
|---|
| 13 | forall(otype K | Comparable(K), otype V)
 | 
|---|
| 14 | int height(tree(K, V) * t){
 | 
|---|
| 15 |   int helper(tree(K, V) * t, int ht){
 | 
|---|
| 16 |     if (empty(t)){
 | 
|---|
| 17 |       return ht;
 | 
|---|
| 18 |     } else if (t->balance > 0){
 | 
|---|
| 19 |       return helper(t->right, 1+ht);
 | 
|---|
| 20 |     } else {
 | 
|---|
| 21 |       // can traverse either branch to find the height
 | 
|---|
| 22 |       // of an AVL tree whose balance is 0
 | 
|---|
| 23 |       return helper(t->left, 1+ht);
 | 
|---|
| 24 |     }
 | 
|---|
| 25 |   }
 | 
|---|
| 26 |   return helper(t, 0);
 | 
|---|
| 27 | }
 | 
|---|
| 28 | 
 | 
|---|
| 29 | forall(otype K | Comparable(K), otype V)
 | 
|---|
| 30 | int calcBalance(tree(K, V) * t){
 | 
|---|
| 31 |   int l = height(t->left);
 | 
|---|
| 32 |   int r = height(t->right);
 | 
|---|
| 33 |   t->balance = r-l;
 | 
|---|
| 34 |   return t->balance;
 | 
|---|
| 35 | }
 | 
|---|
| 36 | 
 | 
|---|
| 37 | // re-establish the link between parent and child
 | 
|---|
| 38 | forall(otype K | Comparable(K), otype V)
 | 
|---|
| 39 | void relinkToParent(tree(K, V) * t){
 | 
|---|
| 40 |   tree(K, V) * parent = t->parent; // FIX ME!!
 | 
|---|
| 41 |   if (empty(t->parent)){
 | 
|---|
| 42 |     return;
 | 
|---|
| 43 |   } else if (parent->key < t->key){
 | 
|---|
| 44 |     parent->right = t;
 | 
|---|
| 45 |   } else {
 | 
|---|
| 46 |     parent->left = t;
 | 
|---|
| 47 |   }
 | 
|---|
| 48 | }
 | 
|---|
| 49 | 
 | 
|---|
| 50 | // rotate left from t
 | 
|---|
| 51 | forall(otype K | Comparable(K), otype V)
 | 
|---|
| 52 | tree(K, V) * rotateLeft(tree(K, V) * t){
 | 
|---|
| 53 |   tree(K, V) * newRoot = t->right;
 | 
|---|
| 54 |   t->right = newRoot->left;
 | 
|---|
| 55 |   newRoot->left = t;
 | 
|---|
| 56 | 
 | 
|---|
| 57 |   // swap parents
 | 
|---|
| 58 |   newRoot->parent = t->parent;
 | 
|---|
| 59 |   t->parent = newRoot;
 | 
|---|
| 60 |   if (t->right != NULL) {
 | 
|---|
| 61 |     tree(K, V) * right = t->right; // FIX ME!!
 | 
|---|
| 62 |     right->parent = t;
 | 
|---|
| 63 |   }
 | 
|---|
| 64 |   // re-establish the link between newRoot and its parent
 | 
|---|
| 65 |   relinkToParent(newRoot);
 | 
|---|
| 66 |   return newRoot;
 | 
|---|
| 67 | }
 | 
|---|
| 68 | 
 | 
|---|
| 69 | // rotate right from t
 | 
|---|
| 70 | forall(otype K | Comparable(K), otype V)
 | 
|---|
| 71 | tree(K, V) * rotateRight(tree(K, V) * t){
 | 
|---|
| 72 |   tree(K, V) * newRoot = t->left;
 | 
|---|
| 73 |   t->left = newRoot->right;
 | 
|---|
| 74 |   newRoot->right = t;
 | 
|---|
| 75 | 
 | 
|---|
| 76 |   // swap parents
 | 
|---|
| 77 |   newRoot->parent = t->parent;
 | 
|---|
| 78 |   t->parent = newRoot;
 | 
|---|
| 79 |   if (t->left != NULL){
 | 
|---|
| 80 |     tree(K, V) * left = t->left; // FIX ME!!
 | 
|---|
| 81 |     left->parent = t;
 | 
|---|
| 82 |   }
 | 
|---|
| 83 |   // re-establish the link between newRoot and its parent
 | 
|---|
| 84 |   relinkToParent(newRoot);
 | 
|---|
| 85 |   return newRoot;
 | 
|---|
| 86 | }
 | 
|---|
| 87 | 
 | 
|---|
| 88 | // balances a node that has balance factor -2 or 2
 | 
|---|
| 89 | forall(otype K | Comparable(K), otype V)
 | 
|---|
| 90 | tree(K, V) * fix(tree(K, V) * t){
 | 
|---|
| 91 |   // ensure that t's balance factor is one of
 | 
|---|
| 92 |   // the appropriate values
 | 
|---|
| 93 |   assert(t->balance == 2 || t->balance == -2);
 | 
|---|
| 94 | 
 | 
|---|
| 95 |   if (t->balance == -2){
 | 
|---|
| 96 |     tree(K, V) * left = t->left; // FIX ME!!
 | 
|---|
| 97 |     if (left->balance == 1){
 | 
|---|
| 98 |       t->left = rotateLeft(t->left);
 | 
|---|
| 99 |     }
 | 
|---|
| 100 |     return rotateRight(t);
 | 
|---|
| 101 |   } else if (t->balance == 2){
 | 
|---|
| 102 |     tree(K, V) * right = t->right; // FIX ME!!
 | 
|---|
| 103 |     if (right->balance == -1){
 | 
|---|
| 104 |       t->right = rotateRight(t->right);
 | 
|---|
| 105 |     }
 | 
|---|
| 106 |     return rotateLeft(t);
 | 
|---|
| 107 |   } else {
 | 
|---|
| 108 |     // shouldn't ever get here
 | 
|---|
| 109 |     assert((int)0);
 | 
|---|
| 110 |     return t;
 | 
|---|
| 111 |   }
 | 
|---|
| 112 | }
 | 
|---|
| 113 | 
 | 
|---|
| 114 | // attempt to fix the tree, if necessary
 | 
|---|
| 115 | forall(otype K | Comparable(K), otype V)
 | 
|---|
| 116 | tree(K, V) * tryFix(tree(K, V) * t){
 | 
|---|
| 117 |   int b = calcBalance(t);
 | 
|---|
| 118 | 
 | 
|---|
| 119 |   if (b == -2 || b == 2){
 | 
|---|
| 120 |     t = fix(t);
 | 
|---|
| 121 |   } else {
 | 
|---|
| 122 |     assert(b == 0 || b == 1 || b == -1);
 | 
|---|
| 123 |   }
 | 
|---|
| 124 |   return t;
 | 
|---|
| 125 | }
 | 
|---|
| 126 | 
 | 
|---|
| 127 | // sets parent field of c to be p
 | 
|---|
| 128 | forall(otype K | Comparable(K), otype V)
 | 
|---|
| 129 | void setParent(tree(K, V) * c, tree(K, V) * p){
 | 
|---|
| 130 |   if (! empty(c)){
 | 
|---|
| 131 |     c->parent = p;
 | 
|---|
| 132 |   }
 | 
|---|
| 133 | }
 | 
|---|
| 134 | 
 | 
|---|