Red Black Tree

redblacktree.hpp

template <typename Key, typename Value>
class RedBlackTree
{
	struct Node;
	Node *nil = new Node();
	Node *root = nil;

	// node structure
	struct Node {
		enum Color { RED, BLACK };
		Color _color;
		Key _key;
		Value _value;
		Node *_parent, *_left, *_right;
		Node(): _color(BLACK) {}
		Node(const Key &key, const Value &value, Node *left, Node *right, Node *parent):
			_key(key), _value(value), _color(RED), _left(left), _right(right), _parent(parent) {}
	};

	// rotation
	void leftRotate(Node *x)
	{
		Node *y = x->_right;
		// remove y->left to x->right
		x->_right = y->_left;
		if (x->_right != nil)
			x->_right->_parent = x;
		// remove y up
		y->_parent = x->_parent;
		if (x->_parent == nil)
			root = y;
		else if (x->_parent->_left == x)
			x->_parent->_left = y;
		else
			x->_parent->_right = y;
		// remove x down
		x->_parent = y;
		y->_left = x;
	}

	void rightRotate(Node *x)
	{
		Node *y = x->_left;
		// remove y->right to x->left
		x->_left = y->_right;
		if (x->_left != nil)
			x->_left->_parent = x;
		// remove y up
		y->_parent = x->_parent;
		if (x->_parent == nil)
			root = y;
		else if (x->_parent->_left == x)
			x->_parent->_left = y;
		else
			x->_parent->_right = y;
		// remove x down
		x->_parent = y;
		y->_right = x;
	}

	Node *min(Node *ptr)
	{
		Node *it = ptr;
		while (it != nil) {
			ptr = it;
			it = it->_left;
		}
		return ptr;
	}

	// insert
	void insertFixup(Node *ptr)
	{
		while (ptr->_parent->_color == Node::RED) {
			if (ptr->_parent == ptr->_parent->_parent->_left) {
				Node *y = ptr->_parent->_parent->_right;
				// case 1
				if (y->_color == Node::RED) {
					ptr->_parent->_color = Node::BLACK;
					y->_color = Node::BLACK;
					ptr->_parent->_parent->_color = Node::RED;
					ptr = ptr->_parent->_parent;
				} else {
					// case 2: switch case 2 to case 3
					if (ptr == ptr->_parent->_right) {
						ptr = ptr->_parent;
						leftRotate(ptr);
					}
					// case 3
					ptr->_parent->_color = Node::BLACK;
					ptr->_parent->_parent->_color = Node::RED;
					rightRotate(ptr->_parent->_parent);
				}
			} else {
				// with 'left' and 'right' exchanged
				Node *y = ptr->_parent->_parent->_left;
				if (y->_color == Node::RED) {
					ptr->_parent->_color = Node::BLACK;
					y->_color = Node::BLACK;
					ptr->_parent->_parent->_color = Node::RED;
					ptr = ptr->_parent->_parent;
				} else {
					if (ptr == ptr->_parent->_left) {
						ptr = ptr->_parent;
						rightRotate(ptr);
					}
					ptr->_parent->_color = Node::BLACK;
					ptr->_parent->_parent->_color = Node::RED;
					leftRotate(ptr->_parent->_parent);
				}
			}
		}
		root->_color = Node::BLACK;
	}

	void insert(Node *nptr)
	{
		Node *it = root, *p = root;
		// find insert position
		while (it != nil) {
			p = it;
			if (nptr->_key < it->_key)
				it = it->_left;
			else if (nptr->_key > it->_key)
				it = it->_right;
			else {
				// find target key-value
				it->_value = nptr->_value;
				return;
			}
		}
		// insert
		nptr->_parent = p;
		if (p == nil)
			root = nptr;
		else if (nptr->_key < p->_key)
			p->_left = nptr;
		else
			p->_right = nptr;
		// fixup
		insertFixup(nptr);
	}

	// find
	Node *find(Key key)
	{
		Node *it = root;
		while (it != nil) {
			if (key < it->_key)
				it = it->_left;
			else if (key > it->_key)
				it = it->_right;
			else
				return it;
		}
		return nil;
	}

	// delete
	void transplant(Node *u, Node *v)
	{
		if (u->_parent == nil)
			root = v;
		else if (u == u->_parent->_left)
			u->_parent->_left = v;
		else
			u->_parent->_right = v;
		v->_parent = u->_parent;
	}

	void removeFixup(Node *ptr)
	{
		while (ptr != root && ptr->_color == Node::BLACK) {
			if (ptr == ptr->_parent->_left) {
				Node *w = ptr->_parent->_right;
				// case 1
				if (w->_color == Node::RED) {
					w->_color = Node::BLACK;
					ptr->_parent->_color = Node::RED;
					leftRotate(ptr->_parent);
					w = ptr->_parent->_right;
				}
				// case 2
				if (w->_left->_color == Node::BLACK && w->_right->_color == Node::BLACK) {
					w->_color = Node::RED;
					ptr = ptr->_parent;
				} else {
					// case 3
					if (w->_right->_color == Node::BLACK) {
						w->_left->_color = Node::BLACK;
						w->_color = Node::RED;
						rightRotate(w);
						w = ptr->_parent->_right;
					}
					// case 4
					w->_color = ptr->_parent->_color;
					ptr->_parent->_color = Node::BLACK;
					w->_right->_color = Node::BLACK;
					leftRotate(ptr->_parent);
					ptr = root;
				}
			} else {
				// with 'left' and 'right' exchanged
				Node *w = ptr->_parent->_left;
				if (w->_color == Node::RED) {
					w->_color = ptr->_parent->_color;
					ptr->_parent->_color = Node::RED;
					rightRotate(ptr->_parent);
					w = ptr->_parent->_left;
				}
				if (w->_left->_color == Node::BLACK && w->_right->_color == Node::BLACK) {
					w->_color = Node::RED;
					ptr = ptr->_parent;
				} else {
					if (w->_left->_color == Node::BLACK) {
						w->_color = Node::RED;
						w->_right->_color = Node::BLACK;
						leftRotate(w);
						w = ptr->_parent->_left;
					}
					w->_color = ptr->_parent->_color;
					w->_left->_color = Node::BLACK;
					ptr->_parent->_color = Node::BLACK;
					rightRotate(ptr->_parent);
					ptr = root;
				}
			}
		}
		ptr->_color = Node::BLACK;
	}

	void remove(Node *ptr)
	{
		Node *y = ptr, *x;
		int y_original_color = y->_color;
		if (y->_left == nil) {
			x = ptr->_right;
			transplant(ptr, ptr->_right);
		} else if (y->_right == nil) {
			x = ptr->_left;
			transplant(ptr, ptr->_left);
		} else {
			y = min(ptr->_right);
			y_original_color = y->_color;
			x = y->_right;
			if (y->_parent == ptr)
				x->_parent = y;	// change nil->_parent
			else {
				transplant(y, y->_right);
				y->_right = ptr->_right;
				y->_right->_parent = y;
			}
			transplant(ptr, y);
			y->_left = ptr->_left;
			y->_left->_parent = y;
			y->_color = ptr->_color;
		}
		if (y_original_color == Node::BLACK)
			removeFixup(x);
	}

	void destruct(Node *node)
	{
		if (node == nil)
			return;
		if (node->_left && node->_left != nil) 
			destruct(node->_left);
		if (node->_right && node->_right != nil)
			destruct(node->_right);
		delete node;
	}

	Node *copy(Node *node, Node *oldNil)
	{
		if (node == oldNil)
			return nil;
		Node *newNode = new Node(node->_key, node->_value, nil, nil, nil);
		newNode->_color = node->_color;
		if (node->_left != oldNil) {
			newNode->_left = copy(node->_left, oldNil);
			newNode->_left->_parent = newNode;
		}
		if (node->_right != oldNil) {
			newNode->_right = copy(node->_right, oldNil);
			newNode->_right->_parent = newNode;
		}
		return newNode;
	}
public:

	RedBlackTree() = default;

	RedBlackTree(const RedBlackTree &tree)
	{
		root = nil = new Node();
		root = copy(tree.root, tree.nil);
	}

	~RedBlackTree()
	{
		destruct(root);
		delete nil;
	}

	RedBlackTree& operator=(RedBlackTree tree)
	{
		std::swap(nil, tree.nil);
		std::swap(root, tree.root);
	}

	Value *get(const Key &key)
	{
		Node *it = find(key);
		return it == nil ? nullptr : &it->_value;
	}

	void put(const Key &key, const Value &value)
	{
		insert(new Node(key, value, nil, nil, nil));
	}

	void remove(const Key &key)
	{
		Node *it = find(key);
		if (it != nil)
			remove(it);
	}
};