平衡二叉树

平衡二叉树的基本操作

#include <cstdio>
#include <algorithm>

using namespace std;

struct Node {
    int data, height;//data为权值，height为当前子树高度
    Node *lchild, *rchild;
};

Node *newNode(int data) {
    Node *node = new Node;
    node->data = data;
    node->height = 1;
    node->lchild = node->rchild = NULL;
    return node;
};

//获取以node为根结点的子树的当前高度
int getHeight(Node *node) {
    if (node == NULL) {
        return 0;
    }
    return node->height;
}

//计算结点node平衡因子
int getBalanceFactor(Node *node) {
    return getHeight(node->lchild) - getHeight(node->rchild);
}

//更新结点的高度
void updateHeight(Node *node) {
    node->height = max(getHeight(node->lchild), getHeight(node->rchild)) + 1;
}

//左旋
void leftRotation(Node *&node) {
    Node *temp = node->rchild;
    node->rchild = temp->lchild;
    temp->lchild = node;
    updateHeight(node);
    updateHeight(temp);
    node = temp;
}

//右旋
void rightRotation(Node *&node) {
    Node *temp = node->lchild;
    node->lchild = temp->rchild;
    temp->rchild = node;
    updateHeight(node);
    updateHeight(temp);
    node = temp;
}

void search(Node *node, int x) {
    if (node == NULL) {
        printf("search failed\n");
        return;
    }
    if (x == node->data) {
        printf("%d\n", node->data);
    } else if (x < node->data) {
        search(node->lchild, x);
    } else {
        search(node->rchild, x);
    }
}

//插入权值为data的结点
void insert(Node *&node, int data) {
    if (node == NULL) {
        node = newNode(data);
        return;
    }
    if (data < node->data) {
        insert(node->lchild, data);
        updateHeight(node);
        if (getBalanceFactor(node) == 2) {
            if (getBalanceFactor(node->lchild) == 1) {//LL型
                rightRotation(node);
            } else if (getBalanceFactor(node->lchild) == -1) {//LR型
                leftRotation(node->lchild);
                rightRotation(node);
            }
        }
    } else {
        insert(node->rchild, data);
        updateHeight(node);
        if (getBalanceFactor(node) == -2) {
            if (getBalanceFactor(node->rchild) == -1) {//RR型
                leftRotation(node);
            } else if (getBalanceFactor(node->rchild) == 1) {//RL型
                rightRotation(node->rchild);
                leftRotation(node);
            }
        }
    }
}

//AVL树的建立
Node *create(int data[], int n) {
    Node *root = NULL;
    for (int i = 0; i < n; i++) {
        insert(root, data[i]);
    }
    return root;
}

1066 Root of AVL Tree