Implementation of a Splay Tree

splay-tree.cpp

#include <iostream>

/*

	Implementation of a splay tree
	Derived from
	 - github.com/jonnyhsy/splay-tree
	 - en.wikipedia.org/wiki/Splay_tree

*/

struct Node {

	Node * left;
	Node * right;
	Node * parent;

	int key;
	int value;

	Node() {
		left = nullptr;
		right = nullptr;
		parent = nullptr;
		key = 0;
		value = 0;
	}

	Node(int k, int v) {
		left = nullptr;
		right = nullptr;
		parent = nullptr;
		key = k;
		value = v;
	}

};

bool isLeftChild(Node * x) {
	return x && x == x -> parent -> left;
}

bool isRightChild(Node * x) {
	return x && x == x -> parent -> right;
}

void printTree(Node * node, int depth = 1) {
	if (node != nullptr) {
		printTree(node -> left, depth + 1);
		for (int i = 0; i < depth; ++i) std::cout << "    ";
		std::cout << node -> key << '\n';
		printTree(node -> right, depth + 1);
	}
}

void recalculate(Node * n) {
	if (n) {
		// do things here...
	}
}

void emplaceChild(Node * x, Node * y) {
	if (x -> parent) {
		if (isLeftChild(x)) {
			x -> parent -> left = y;
		} else {
			x -> parent -> right = y;
		}
	}
}

/*
	
	Right Rotation

        |                  |
        X                  Y 
       / \                / \
      Y  C     ==>       A  X 
     / \                   / \
    A  B                  B  C

*/

void rotateRight(Node * x) {
	Node * y = x -> left;
	if (y) {
		emplaceChild(x, y);
		if (y -> right) y -> right -> parent = x;
		x -> left = y -> right;
		y -> right = x;
		y -> parent = x -> parent;
		x -> parent = y;
		recalculate(x);
		recalculate(y);
		recalculate(y -> parent);
	}
}

/*

	Left Rotation

        |                       |
        x                       y
       / \                     / \
      A  y         ==>        x  C
        / \                  / \  
       B  C                 A  B  

*/

void rotateLeft(Node * x) {
	Node * y = x -> right;
	if (y) {
		emplaceChild(x, y);
		if (y -> left) y -> left -> parent = x;
		x -> right = y -> left;
		y -> left = x;
		y -> parent = x -> parent;
		x -> parent = y;
		recalculate(x);
		recalculate(y);
		recalculate(y -> parent);
	}
}

Node * splay(Node * x) {
	while (x -> parent) {
		if (x -> parent -> parent == nullptr) {
			if (isLeftChild(x)) rotateRight(x -> parent);
			else rotateLeft(x -> parent);
		} else if (isLeftChild(x) && isLeftChild(x -> parent)) {
			rotateRight(x -> parent -> parent);
			rotateRight(x -> parent);
		} else if (isRightChild(x) && isRightChild(x -> parent)) {
			rotateLeft(x -> parent -> parent);
			rotateLeft(x -> parent);
		} else if (isLeftChild(x) && isRightChild(x -> parent)) {
			rotateRight(x -> parent);
			rotateLeft(x -> parent);
		} else {
			rotateLeft(x -> parent);
			rotateRight(x -> parent);
		}
	}
	return x;
}

Node * find(Node * n, int k) {
	while (n && n -> key != k) n = (k < n -> key) ? n -> left : n -> right;
	if (n) splay(n);
	return n;
}

Node * largest(Node * n) {
	while (n -> right) n = n -> right;
	return n;
}

Node * remove(Node * n) {
	splay(n);
	Node * a = n -> left;
	Node * b = n -> right;
	Node * c = a ? a : b;
	if (a) a -> parent = nullptr;
	if (b) b -> parent = nullptr;
	if (a && b) {
		c = splay(largest(a));
		c -> right = b;
		b -> parent = c;
		recalculate(c);
	}
	return c;
}

Node * insert(Node * root, int value, int key) {
	Node * newnode = new Node(value, key);
	Node * current = root;
	Node * last = nullptr;
	bool direction = false;
	while (current != nullptr) {
		last = current;
		if (key > current -> key) {
			current = current -> right;
			direction = true;
		} else {
			current = current -> left;
			direction = false;
		}
	}
	if (direction) last -> right = newnode;
	else last -> left = newnode;
	newnode -> parent = last;
	splay(newnode);
	return newnode;
}

int main() {
	std::cout << "DXsmiley's splay tree implementation." << std::endl;
	std::cout << "Operations are Splay, Insert, Remove." << std::endl;
	Node * root = new Node(0, 0);
	while (true) {
		std::cout << "---------------------------------------" << std::endl;
		printTree(root);	
		std::cout << "---------------------------------------" << std::endl;
		std::string command;
		int key;
		std::cin >> command >> key;
		if (command == "s" || command == "S") root = splay(find(root, key));
		if (command == "i" || command == "I") root = insert(root, key, key);
		if (command == "r" || command == "R") root = remove(find(root, key));
	}
}