1 /*
  2     平衡二叉树
  3 */
  4
  5 #include <stdio.h>
  6 #include <string.h>
  7 #include <math.h>
  8 #include <stdlib.h>
  9 #include <time.h>
 10 #include <stdbool.h>
 11
 12 // 平衡二叉树结构体
 13 typedef struct _node {
 14     int v, height;            // v 为结点权值， height 为当前子树高度
 15     struct _node *lchild, *rchild;    // 左右孩子结点地址
 16 } node;
 17
 18 // 平衡二叉树新建结点
 19 node* newNode(int v) {
 20     node* Node = (node*)malloc(sizeof(node));    // 申请地址空间
 21     Node->v = v;                                // 结点权值为 v
 22     Node->height = 1;                            // 结点高度初始为 1
 23     Node->lchild = Node->rchild = NULL;            // 初始状态下没有左右孩子
 24     return Node;
 25 }
 26
 27 // 获取以 root 为根结点的子树的当前高度
 28 int getHeight(node* root) {
 29     if(root == NULL)    return 0;    // 空树高度为 0
 30     return root->height;
 31 }
 32
 33 // 计算结点 root 的平衡因子
 34 int getBalanceFactor(node* root) {
 35     // 左子树高度减右子树高度
 36     return getHeight(root->lchild) - getHeight(root->rchild);
 37 }
 38
 39 // 更新结点 root 的高度
 40 void updateHeight(node* root) {
 41     // 根结点高度为左右子树高度较大值 +1
 42     int lHeight = getHeight(root->lchild), rHeight = getHeight(root->rchild);
 43     int max = lHeight > rHeight ? lHeight : rHeight;
 44     root->height = max + 1;
 45 }
 46
 47 // 查找平衡二叉树中数据域为 x 的结点
 48 void search(node* root, int x) {
 49     if(root == NULL) {    // 空树，查找失败
 50         printf("search failed\n");
 51         return;
 52     }
 53     if(x == root->v) {    // 查找成功，访问之
 54         printf("search success %d\n", root->v);
 55     } else if(x < root->v) {        // x 比根结点的数据域小，往左子树查找
 56         search(root->lchild, x);
 57     } else {
 58         search(root->rchild, x);    // x 比根结点的数据域大，往右子树查找
 59     }
 60 }
 61
 62 // 左旋
 63 void L(node** root) {
 64     node* temp = (*root)->rchild;    // root指向A,temp指向B
 65     (*root)->rchild = temp->lchild;    // 步骤 1
 66     temp->lchild = *root;            // 步骤 2
 67     updateHeight(*root);                // 更新高度
 68     updateHeight(temp);
 69     (*root) = temp;                    // 步骤 3
 70 }
 71
 72 // 右旋
 73 void R(node** root) {
 74     node* temp = (*root)->lchild;    // root指向A,temp指向B
 75     (*root)->lchild = temp->rchild;    // 步骤 1
 76     temp->rchild = *root;            // 步骤 2
 77     updateHeight(*root);                // 更新高度
 78     updateHeight(temp);
 79     (*root) = temp;                    // 步骤 3
 80 }
 81
 82 // 插入权值为 v 的结点
 83 void insert(node** root, int v) {
 84     if((*root) == NULL) {        // 到达插入位置
 85         (*root) = newNode(v);
 86         return;
 87     }
 88     if(v < (*root)->v) {        // v 比根结点权值小
 89         insert(&(*root)->lchild, v);        // 往左子树插入
 90         updateHeight(*root);                // 更新树高
 91         if(getBalanceFactor(*root) == 2) {
 92             if(getBalanceFactor((*root)->lchild) == 1) {    // LL 型
 93                 R(*root);
 94             } else if(getBalanceFactor((*root)->lchild) == -1) {    // LR 型
 95                 L(&(*root)->lchild);
 96                 R(root);
 97             }
 98         }
 99     } else {                    // v 比根结点权值大
100         insert(&(*root)->rchild, v);        // 往右子树插入
101         updateHeight(*root);                // 更新树高
102         if(getBalanceFactor(*root) == -2) {
103             if(getBalanceFactor((*root)->rchild) == -1) {    // RR 型
104                 L(root);
105             } else if(getBalanceFactor((*root)->rchild) == 1) {    // RL 型
106                 R(&(*root)->rchild);
107                 L(root);
108             }
109         }
110     }
111 }
112
113 // AVL 树的建立
114 node* create(int data[], int n) {
115     node* root = NULL;        // 新建根结点 root
116     int i;
117     for(i=0; i<n; ++i) {
118         insert(&root, data[i]);        // 将 data 中数据依次插入AVL树
119     }
120     return root;            // 返回根结点
121 }
122
123 // 先序遍历
124 void preorder(node* root) {
125     if(root == NULL) {
126         return;        // 空树，递归边界
127     }
128     printf("%d\n", root->v);        // 访问该结点
129     preorder(root->lchild);            // 访问左子树
130     preorder(root->rchild);            // 访问右子树
131 }
132
133 int main() {
134     int data[10] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
135     node* root = NULL;
136     root = create(data, 10);    // 创建 AVL 树
137     preorder(root);                // 先序遍历
138     search(root, 10);            // 在 root 中查找 10
139     return 0;
140 }