AVL树中,当进行分割操作时,需要计算新生成的两棵子树的大小,即左子树的大小和右子树的大小。下面是一个示例代码:
struct Node { int val; int height; int size; Node *left; Node *right;
Node(int val) { this->val = val; height = 1; size = 1; left = NULL; right = NULL; } };
int getSize(Node *root) { if (!root) { return 0; } return root->size; }
void updateSize(Node *root) { root->size = getSize(root->left) + getSize(root->right) + 1; }
void updateHeight(Node *root) { root->height = max(getHeight(root->left), getHeight(root->right)) + 1; }
void update(Node *root) { updateSize(root); updateHeight(root); }
int getBalanceFactor(Node *root) { if (!root) { return 0; } return getHeight(root->left) - getHeight(root->right); }
Node *rightRotate(Node *root) { Node *newRoot = root->left; root->left = newRoot->right; newRoot->right = root; update(root); update(newRoot); return newRoot; }
Node *leftRotate(Node *root) { Node *newRoot = root->right; root->right = newRoot->left; newRoot->left = root; update(root); update(newRoot); return newRoot; }
Node *insert(Node *root, int val) { if (!root) { return new Node(val); }
if (val < root->val) { root->left = insert(root->left, val); } else { root->right = insert(root->right, val); }
update(root);
int balanceFactor = getBalanceFactor(root);
if (balanceFactor > 1 && getBalanceFactor(root->left) >= 0) { return rightRotate(root); } if (balanceFactor < -1