Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include <iostream>
- #include <algorithm>
- struct Node {
- int value;
- int height;
- struct Node *left;
- struct Node *right;
- };
- class AVLTree {
- private:
- Node *root_;
- int size_;
- void clear(Node *n) {
- if (n == nullptr) {
- return;
- }
- clear(n->left);
- clear(n->right);
- delete n;
- }
- int height(Node *n) {
- if (n == nullptr) {
- return 0;
- }
- return n->height;
- }
- void traversal(Node *n, int *p, int *i) {
- if (n == nullptr) {
- return;
- }
- traversal(n->left, p, i);
- *(p + *i) = n->value;
- (*i)++;
- traversal(n->right, p, i);
- }
- int balance(Node *n) {
- if (n == nullptr) {
- return 0;
- }
- return height(n->left) - height(n->right);
- }
- Node *insert(int val, Node *n) {
- if (n == nullptr) {
- n = new Node;
- n->value = val;
- n->left = nullptr;
- n->right = nullptr;
- n->height = 1;
- size_++;
- return n;
- }
- if (val < n->value) {
- n->left = insert(val, n->left);
- } else if (val > n->value) {
- n->right = insert(val, n->right);
- } else {
- return n;
- }
- updateHeight(n);
- return rebalance(n);
- }
- Node *erase(Node *n, int val) {
- if (n == nullptr) {
- return nullptr;
- }
- if (n->value > val) {
- n->left = erase(n->left, val);
- } else if (n->value < val) {
- n->right = erase(n->right, val);
- } else {
- Node *left = n->left;
- Node *right = n->right;
- delete n;
- size_--;
- if (left == nullptr) {
- return right;
- } else if (right == nullptr) {
- return left;
- }
- Node *tmp = right;
- while (tmp->left != nullptr) {
- tmp = tmp->left;
- }
- tmp->left = left;
- updateHeight(tmp);
- return rebalance(tmp);
- }
- updateHeight(n);
- return rebalance(n);
- }
- Node *rebalance(Node *n) {
- int b = balance(n);
- if (b > 1 && balance(n->left) < 0) {
- n->left = leftRotate(n->left);
- return rightRotate(n);
- } else if (b < -1 && balance(n->right) > 0) {
- n->right = rightRotate(n->right);
- return leftRotate(n);
- } else if (b > 1) {
- return rightRotate(n);
- } else if (b < -1) {
- return leftRotate(n);
- }
- return n;
- }
- Node *rightRotate(Node *n) {
- Node *rotated = n->left;
- Node *right = rotated->right;
- rotated->right = n;
- n->left = right;
- updateHeight(n);
- updateHeight(rotated);
- return rotated;
- }
- Node *leftRotate(Node *n) {
- Node *rotated = n->right;
- Node *left = rotated->left;
- rotated->left = n;
- n->right = left;
- updateHeight(n);
- updateHeight(rotated);
- return rotated;
- }
- void updateHeight(Node *n) {
- int h = std::max(height(n->left), height(n->right));
- n->height = h + 1;
- }
- Node *find(Node *n, int val) {
- if (n == nullptr) {
- return nullptr;
- }
- if (n->value > val) {
- return find(n->left, val);
- } else if (n->value < val) {
- return find(n->right, val);
- } else {
- return n;
- }
- }
- public:
- AVLTree() {
- root_ = nullptr;
- size_ = 0;
- }
- int getHeight() {
- return height(root_);
- }
- Node *getRoot() {
- return root_;
- }
- bool empty() {
- return (root_ == nullptr);
- }
- int getSize() {
- return size_;
- }
- int *find(int value) {
- Node *res = find(root_, value);
- if (res != nullptr) {
- return &(res->value);
- }
- return nullptr;
- }
- void insert(int val) {
- root_ = insert(val, root_);
- }
- void erase(int value) {
- root_ = erase(root_, value);
- }
- int *traversal() {
- if (root_ == nullptr) {
- return nullptr;
- }
- int *arr = new int[getSize()];
- int i = 0;
- traversal(root_, arr, &i);
- return arr;
- }
- int *lowerBound(int val) {
- int *arr = traversal();
- int *ans = nullptr;
- int size = getSize();
- for (int i = 0; i < size; ++i) {
- if (arr[i] >= val) {
- ans = find(arr[i]);
- break;
- }
- }
- delete arr;
- return ans;
- }
- ~AVLTree() {
- clear(root_);
- }
- };
- int main() {
- AVLTree tree;
- std::cout << tree.lowerBound(9) << std::endl;
- tree.erase(1000);
- tree.insert(0);
- tree.insert(1);
- tree.insert(2);
- std::cout << tree.lowerBound(0) << std::endl;
- tree.erase(1);
- tree.erase(2);
- tree.erase(0);
- std::cout << tree.lowerBound(0) << std::endl;
- return 0;
- }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement