PTA (Advanced Level)1066 Root of AVL Tree

  1 #include <bits/stdc++.h>
  2 using namespace std;
  3 typedef int dataType;
  4 vector<dataType> data;
  5 struct node{
  6     dataType data;
  7     int height; //AVL树结点比起普通二叉搜索树需要记录height
  8     node *leftChild;
  9     node * rightChild;
 10     node(){
 11         height = 1;
 12         leftChild = NULL;
 13         rightChild = NULL;
 14     }
 15 };
 16 int getHeight(node *root){  //获取高度
 17     if(root == NULL)  
 18         return 0;
 19     else
 20         return root->height;
 21 }
 22 int getBalanceFactor(node *root){   //获取树的叶子结点高度差左高为正右高为负
 23     return getHeight(root->leftChild) - getHeight(root->rightChild);
 24 }
 25 int updateHeight(node *root){   //更新高度
 26     root->height = max(getHeight(root->leftChild),getHeight(root->rightChild)) + 1;
 27 }
 28 void leftRotation(node *&root){ //左旋
 29     node *temp = root->rightChild;      //root指向先前根结点temp指向右子树根结点
 30     root->rightChild = temp->leftChild; //temp指向根结点的右子树，所以其所有结点都大于根结点
 31     //由于在左旋中需要使temp成为新的根结点，所以将root右子树指向temp左子树，再让temp左子树指向root
 32     temp->leftChild = root;
 33     //更新root与temp的树高
 34     updateHeight(root);
 35     updateHeight(temp);
 36     root = temp;    //temp成为新的根结点
 37 }
 38 void rightRotation(node *&root){    //右旋思路同左旋
 39     node *temp = root->leftChild;
 40     root->leftChild = temp->rightChild;
 41     temp->rightChild = root;
 42     updateHeight(root);
 43     updateHeight(temp);
 44     root = temp;
 45 }
 46 void insertAVLTree(node *&root, int x){ //插入结点
 47     if(root == NULL){   //找到插入位置
 48         root = new node();
 49         root->data = x;
 50         return;
 51     }
 52     if(root->data == x){    //结点已存在
 53         return;
 54     }else if(root->data > x){   //要插入的数据比根结点权值小
 55         insertAVLTree(root->leftChild, x);  //插入左子树
 56         updateHeight(root);
 57         if(getBalanceFactor(root) == 2){
 58             if(getBalanceFactor(root->leftChild) == 1){
 59                 rightRotation(root);
 60             }else if(getBalanceFactor(root->leftChild) == -1){
 61                 leftRotation(root->leftChild);
 62                 rightRotation(root);
 63             }
 64         }
 65     }else if(root->data < x){   //要插入的数据比根结点权值大
 66         insertAVLTree(root->rightChild, x); //插入右子树
 67         updateHeight(root);
 68         if(getBalanceFactor(root) == -2){
 69             if(getBalanceFactor(root->rightChild) == -1){
 70                 leftRotation(root);
 71             }else if(getBalanceFactor(root->rightChild) == 1){
 72                 rightRotation(root->rightChild);
 73                 leftRotation(root);
 74             }
 75         }
 76     }
 77 }
 78 node *createAVLTree(){
 79     node *root = NULL;
 80     for(vector<dataType>::iterator it = data.begin(); it != data.end(); it++){
 81         insertAVLTree(root, *it);
 82     }
 83     return root;
 84 }
 85 /*void preorder(node *root){
 86     if(root == NULL)
 87         return;
 88     cout << root -> data << " ";
 89     preorder(root -> leftChild);
 90     preorder(root -> rightChild);
 91 }*/
 92 int main()
 93 {
 94     int n;
 95     while(scanf("%d", &n) != EOF){
 96         data.clear();
 97         for(int i = 0; i < n; i++){
 98             dataType temp;
 99             scanf("%d", &temp);
100             data.push_back(temp);
101         }
102         node *root = createAVLTree();
103         //preorder(root);
104         //cout << endl;
105         printf("%d\n", root->data);
106     }
107     return 0;
108 }