RB Tree

  1 //============================================================================
  2 // Name        : Tree.cpp
  3 // Author      : 
  4 // Version     :
  5 // Copyright   : Your copyright notice
  6 // Description : Hello World in C++, Ansi-style
  7 //============================================================================
  8 
  9 #include <cmath>
 10 #include <iostream>
 11 
 12 using namespace std;
 13 
 14 #define RED        true
 15 #define BLACK    false
 16 
 17 template <typename T>
 18 class Node{
 19 public:
 20     T data;
 21     bool color;
 22     Node<T> *L;
 23     Node<T> *R;
 24     Node<T> *P;
 25     Node():data(0), color(RED), L(NULL), R(NULL), P(NULL){}
 26 };
 27 
 28 template <typename T>
 29 class AVLTree{
 30 private:
 31     Node<T> *insert(Node<T> *S, T data);
 32     Node<T> *adj(Node<T> *S);
 33     Node<T> *singleRorateLeft(Node<T> *S);
 34     Node<T> *singleRorateRight(Node<T> *S);
 35 public:
 36     Node<T> *root;
 37     AVLTree():root(NULL){};
 38 
 39     int depth(Node<T> *S);
 40     Node<T> *check(int pos);
 41 
 42     void add(T data){adj(insert(root, data));}
 43 };
 44 
 45 template <typename T> int AVLTree<T>::depth(Node<T> *S)
 46 {
 47     if(NULL == S)
 48         return 0;
 49     return depth(S->L) > depth(S->R) ? depth(S->L) + 1 : depth(S->R) + 1;
 50 }
 51 
 52 template <typename T> Node<T> *AVLTree<T>::check(int pos)
 53 {
 54     if(pos == 1)
 55         return root;
 56     Node<T> *parent = check(pos/2);
 57     if(NULL != parent)
 58         return pos%2 ? parent->R : parent->L;
 59     else
 60         return NULL;
 61 }
 62 
 63 template <typename T> Node<T> *AVLTree<T>::singleRorateLeft(Node<T> *S)
 64 {
 65     Node<T> *K = S->L;
 66     if(K->R != NULL)
 67         K->R->P = S;
 68     K->P = S->P;
 69     if(S->P != NULL)
 70         if(S->P->L == S)
 71             S->P->L = K;
 72         else
 73             S->P->R = K;
 74     S->P = K;
 75     S->L = K->R;
 76     K->R = S;
 77     if(root == S)
 78         root = K;
 79     return K;
 80 }
 81 
 82 template <typename T> Node<T> *AVLTree<T>::singleRorateRight(Node<T> *S)
 83 {
 84     Node<T> *K = S->R;
 85     if(K->L != NULL)
 86         K->L->P = S;
 87     K->P = S->P;
 88     if(S->P != NULL)
 89         if(S->P->L == S)
 90             S->P->L = K;
 91         else
 92             S->P->R = K;
 93     S->P = K;
 94     S->R = K->L;
 95     K->L = S;
 96     if(root == S)
 97         root = K;
 98     return K;
 99 }
100 template <typename T> Node<T> * AVLTree<T>::adj(Node<T> *S)
101 {
102     if(S == root || S->color == BLACK)
103         root->color = BLACK;
104     else if(S->P->color == RED)
105     {
106         Node<T> *U = (S->P->P->L == S->P) ? S->P->P->R : S->P->P->L;
107         if(U != NULL && U->color == RED)
108         {
109             S->P->P->color = RED;
110             U->color = BLACK;
111             S->P->color = BLACK;
112             return adj(S->P->P);
113         }
114         else if(S->P->P->L == S->P)
115         {
116             if(S->P->L == S)
117             {
118                 S->P->P->color = RED;
119                 S->P->color = BLACK;
120                 return singleRorateLeft(S->P->P);
121             }
122             else
123             {
124                 return adj(singleRorateRight(S->P)->L);
125             }
126         }
127         else
128         {
129             if(S->P->L == S)
130             {
131                 return adj(singleRorateLeft(S->P)->R);
132             }
133             else
134             {
135                 S->P->P->color = RED;
136                 S->P->color = BLACK;
137                 return singleRorateRight(S->P->P);
138             }
139         }
140     }
141     else
142         return S;
143 }
144 
145 template <typename T> Node<T> * AVLTree<T>::insert(Node<T> *S, T data)
146 {
147     if(root == NULL)
148     {
149         root = new Node<T>();
150         root->data = data;
151         return root;
152     }
153     else if(data > S->data)
154         if(S->R == NULL)
155         {
156             S->R = new Node<T>();
157             S->R->data = data;
158             S->R->P = S;
159             return S->R;
160         }
161         else
162             return insert(S->R, data);
163     else if(data < S->data)
164         if(S->L == NULL)
165         {
166             S->L = new Node<T>();
167             S->L->data = data;
168             S->L->P = S;
169             return S->L;
170         }
171         else
172             return insert(S->L, data);
173     else
174         return S;
175 }
176 
177 template <typename T> ostream &operator<<(ostream &out,AVLTree<T> &R)
178 {
179     for(int i = 1; i < pow(2, R.depth(R.root)); i++)
180     {
181         if(R.check(i) == NULL)
182             cout<<i<<"\t"<<"NULL"<<endl;
183         else
184             cout<<i<<"\t"<<R.check(i)->data<<"\t"<<R.check(i)->color<<endl;
185     }
186     return out;
187 }
188 
189 int main()
190 {
191     AVLTree<int> t;
192     t.add(1);
193     t.add(2);
194     t.add(3);
195     t.add(4);
196     t.add(5);
197     t.add(6);
198     t.add(7);
199     t.add(8);
200     t.add(9);
201     t.add(10);
202     t.add(11);
203     t.add(12);
204     t.add(13);
205     t.add(14);
206     t.add(15);
207     t.add(16);
208     t.add(17);
209     t.add(18);
210     t.add(19);
211     t.add(20);
212     cout<<t<<endl;
213 }