平衡二叉树

  1 #include <iostream>
  2 #include <cstdio>
  3 #include <cstring>
  4 #include <vector>
  5 
  6 using namespace std;
  7 
  8 template <typename E> class BSTNode{
  9 public:
 10     E data;
 11     int bf;
 12     BSTNode* lch,*rch;
 13     BSTNode() {
 14         lch = rch = NULL;
 15     }
 16 };
 17 
 18 template <typename E> class BalanceTree {
 19 private:
 20     void R_Rotate (BSTNode<E>*& p) {
 21         BSTNode<E>* lc = p -> lch;
 22         p -> lch = lc -> rch;
 23         lc -> rch = p;
 24         p = lc;
 25     }
 26 
 27     void L_Rotate(BSTNode<E>*& p) {
 28         BSTNode<E>* rc = p -> rch;
 29         p -> rch = rc -> lch;
 30         rc -> lch = p;
 31         p = rc;
 32     }
 33 
 34     #define LH 1
 35     #define EH 0
 36     #define RH -1
 37 
 38     void LeftBalace(BSTNode<E>*& T) {
 39         BSTNode<E>* lc = T -> lch;
 40         switch (lc -> bf) {
 41             case LH:
 42                 T -> bf = lc -> bf = EH;
 43                 R_Rotate(T);
 44                 break;
 45             case RH:
 46                 BSTNode<E>* rd = lc -> rch;
 47                 switch (rd -> bf) {
 48                     case LH: T -> bf = RH;lc -> bf = EH;break;
 49                     case EH: T -> bf = lc -> bf = EH; break;
 50                     case RH: T -> bf = EH;lc -> bf = LH;break;
 51                 }
 52                 rd -> bf = EH;
 53                 L_Rotate(lc);
 54                 R_Rotate(T);
 55         }
 56     }
 57 
 58     void RightBalace(BSTNode<E>*& T) {
 59         BSTNode<E>* rc,*rd;
 60         rc = T -> rch;
 61         switch (rc -> bf) {
 62             case RH:
 63                 T -> bf = rc -> bf = EH;
 64                 L_Rotate(T);
 65                 break;
 66             case LH:
 67                 rd = rc -> lch;
 68                 switch (rd -> bf) {
 69                     case RH: T -> bf = LH;rc -> bf = EH;break;
 70                     case EH: T -> bf = rc -> bf = EH; break;
 71                     case LH: T -> bf = EH;rc -> bf = RH;break;
 72                 }
 73             rd -> bf = EH;
 74             R_Rotate(rc);
 75             L_Rotate(T);
 76         }
 77     }
 78 
 79     int InsertElem(BSTNode<E>*& T,E data,int& taller) {
 80         if (!T) {
 81             T = new BSTNode<E>;
 82             T -> data = data;
 83             T -> bf = EH;
 84             taller = 1;
 85         }
 86         else {
 87             if (data == T -> data) {
 88                 taller = 0;
 89                 return 0;
 90             }
 91             if (data < T -> data) {
 92                 if (!InsertElem(T -> lch,data,taller)) return 0;
 93                 if (taller) {
 94                     switch(T -> bf) {
 95                         case LH:
 96                             LeftBalace(T);
 97                             taller = 0;
 98                             break;
 99                         case EH:
100                             T -> bf = LH;
101                             taller = 1;
102                             break;
103                         case RH:
104                             T -> bf = EH;
105                             taller = 0;
106                             break;
107                     }
108                 }
109             }
110             else {
111                 if (!InsertElem(T -> rch,data,taller)) return 0;
112                 if (taller) {
113                     switch(T -> bf) {
114                         case LH:
115                             T -> bf = EH;
116                             taller = 0;
117                             break;
118                         case EH:
119                             T -> bf = RH;
120                             taller = 1;
121                             break;
122                         case RH:
123                             RightBalace(T);
124                             taller = 0;
125                             break;
126                     }
127                 }
128             }
129         }
130         return 1;
131     }
132 
133     BSTNode<E>* SearchTree(BSTNode<E>*& T,E data) {
134         if (T) {
135             if (data == T -> data) return T;
136             else if (data < T -> data) return SearchTree(T -> lch,data);
137             else return SearchTree(T -> rch,data);
138         }
139         else return NULL;
140     }
141 
142 public:
143     BSTNode<E>* head;
144     void insert(E data) {
145         int taller = 0;
146         InsertElem(head,data,taller);
147     }
148     int search(E data) {
149         BSTNode<E>* temp = SearchTree(head,data);
150         if (temp) return 1;
151         return 0;
152     }
153 };
154 
155 int main () {
156     int n;
157     freopen("1.in","r",stdin);
158     BalanceTree<int> T;
159     cin >> n;
160     for (int i =0 ;i < n;i++) {
161         int data;
162         cin >> data;
163         T.insert(data);
164     }
165     int m;
166     cin >> m;
167     for (int i = 0;i < m;i++) {
168         int data;
169         cin >> data;
170         cout << T.search(data) << endl;
171     }
172 }