Skip to content

Instantly share code, notes, and snippets.

@vetom
Last active February 24, 2019 00:49
Show Gist options
  • Save vetom/119761faf92f294619a97a014efa60df to your computer and use it in GitHub Desktop.
Save vetom/119761faf92f294619a97a014efa60df to your computer and use it in GitHub Desktop.
C++ AVL Tree
#include <iostream>
#include <algorithm> // swap, sort, min, find, max
using namespace std;
class Node
{
int data;
Node* left;
Node* right;
int height;
void updateHeight()
{
int leftHeight(0), rightHeight(0);
if (left != nullptr) leftHeight = left->height;
if (right != nullptr) rightHeight = right->height;
height = max(leftHeight, rightHeight) + 1;
}
Node* insertSubTree(Node* subTree)
{
if (left == nullptr) left = subTree;
else left->insertSubTree(subTree);
return this;
}
Node* rotateLL()
{
Node* head = this->left;
this->left = head->right;
head->right = this;
return head;
}
Node* rotateRR()
{
Node* head = this->right;
this->right = head->left;
head->left = this;
return head;
}
Node* rotateRL()
{
Node* head = this->right->left;
this->right->left = head->right;
head->right = this->right;
this->right = head->left;
head->left = this;
return head;
}
Node* rotateLR()
{
Node* head = this->left->right;
this->left->right = head->left;
head->left = this->left;
this->left = head->right;
head->right = this;
return head;
}
Node* balanceSubTree()
{
int r(0), l(0), dr(0), dl(0), dbalance;
if (left != nullptr) l = left->height;
if (right != nullptr) r = right->height;
int balance = l - r;
if (balance > -2 && balance < 2) return this;
if (balance > 0)
{
if (left->left != nullptr) dl = left->left->height;
if (left->right != nullptr) dr = left->right->height;
dbalance = dl - dr;
if (dbalance > 0) return rotateLL();
else return rotateLR();
}
else
{
if (right->left != nullptr) dl = right->left->height;
if (right->right != nullptr) dr = right->right->height;
dbalance = dl - dr;
if (dbalance > 0) return rotateRL();
else return rotateRR();
}
}
public:
Node(int value) : data(value), left(nullptr), right(nullptr), height(1) { }
void updateTreeHeight()
{
if (right != nullptr) right->updateTreeHeight();
if (left != nullptr) left->updateTreeHeight();
updateHeight();
}
Node* balanceTree()
{
if (right != nullptr) right = right->balanceTree();
if (left != nullptr) left = left->balanceTree();
return balanceSubTree();
}
Node* insert(int value)
{
if (value < data)
{
if (left == nullptr) left = new Node(value);
else left->insert(value);
}
else
{
if (right == nullptr) right = new Node(value);
else right->insert(value);
}
updateHeight();
return this;
}
bool check(int value)
{
if (data == value) return true;
if (value < data && left != nullptr) return left->check(value);
else if (right != nullptr) return right->check(value);
return false;
}
void printASC()
{
if (left != nullptr)left->printASC();
cout << data << "(" << height << ")" << " ";
if (right != nullptr)right->printASC();
}
void printDSC()
{
if (right != nullptr)right->printDSC();
cout << data << "(" << height << ")" << " ";
if (left != nullptr)left->printDSC();
}
Node* remove(int value)
{
if (data == value)
{
if (left == nullptr) return right;
else if (right == nullptr) return left;
else return right->insertSubTree(left);;
}
else
{
if (value < data) left = left->remove(value);
else right = right->remove(value);
}
return this;
}
};
void showCommandList()
{
cout << "---- COMMAND LIST ----" << endl;
cout << " I n : INSERT VALUE n " << endl;
cout << " M k [n1, n2, ..., nk] : INSERT k VALUES " << endl;
cout << " R n : REMOVE VALUE n " << endl;
cout << " C n : IS n INSERTED? " << endl;
cout << " A : PRINT ASC WITH HEIGHT " << endl;
cout << " D : PRINT DSC WITH HEIGHT " << endl;
cout << " ? : SHOW COMMAND LIST " << endl;
cout << " E : EXIT PROGRAM " << endl;
}
int main()
{
char command;
int n, k;
Node* tree;
showCommandList();
while (true)
{
cin >> command;
if (command == 'e' || command == 'E') return 0;
if (command == '?') showCommandList();
else if (command == 'i' || command == 'I')
{
cin >> n;
if (tree == nullptr) tree = new Node(n);
else tree->insert(n);
tree = tree->balanceTree();
tree->updateTreeHeight();
}
else if (command == 'm' || command == 'M')
{
cin >> k;
for (int i = 0; i < k; i++)
{
cin >> n;
if (tree == nullptr) tree = new Node(n);
else tree = tree->insert(n);
tree = tree->balanceTree();
tree->updateTreeHeight();
}
}
else if (command == 'c' || command == 'C')
{
cin >> n;
if (tree == nullptr) cout << "NO" << endl;
else if (tree->check(n)) cout << "YES" << endl;
else cout << "NO" << endl;
}
else if (command == 'a' || command == 'A')
{
if (tree == nullptr) cout << "It's empty!";
else tree->printASC();
cout << endl;
}
else if (command == 'd' || command == 'D')
{
if (tree == nullptr) cout << "It's empty!";
else tree->printDSC();
cout << endl;
}
else if (command == 'r' || command == 'R')
{
cin >> n;
if (tree == nullptr) cout << "It's empty!" << endl;
else tree = tree->remove(n);
tree->updateTreeHeight();
tree = tree->balanceTree();
tree->updateTreeHeight();
}
}
return 0;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment