natalie_code/program3/AVLCommands.cpp

// height function partly based on code from
// http://www.sanfoundry.com/cpp-program-implement-avl-trees/
// balance function partly based on code from
// http://www.cplusplus.com/forum/beginner/54835/

#include <cassert>
#include <iostream>
#include <string>
#include <queue>

#include "json.hpp"
#include "AVLCommands.h"

BSTNode::BSTNode(int key) :
	key_(key),
	parent_(std::weak_ptr<BSTNode>()),
	left_(nullptr),
	right_(nullptr),
	height_(0),
	balance_(0) {} // add a height_

BSTNode::BSTNode(int key, std::weak_ptr<BSTNode> parent) :
	key_(key),
	parent_(parent),
	left_(nullptr),
	right_(nullptr),
	height_(0),
	balance_(0) {}

bool BSTNode::IsLeaf() const {
	return left_ == nullptr && right_ == nullptr;
}

bool BSTNode::HasLeftChild() const {
	return left_ != nullptr;
}

bool BSTNode::HasRightChild() const {
	return right_ != nullptr;
}

void BSTNode::DeleteChild(std::shared_ptr<BSTNode> v) {
	if (left_ == v) {
		left_ = nullptr;
	} else if (right_ == v) {
		right_ = nullptr;
	} else {
		std::cerr << "BSTNode::DeleteChild Error: non-child passed as argument\n";
		exit(EXIT_FAILURE);
	}
}

void BSTNode::ReplaceChild(std::shared_ptr<BSTNode> v, std::shared_ptr<BSTNode> u) {
	//std::cout << "v, child: " << v->key_ << std::endl;
	//std::cout << "u, parent: " << u->key_ << std::endl;
	if (left_ == u || right_ == u) {
		std::cerr << "BSTNode::ReplaceChild Error: child passed as replacement\n";
	}
	if (left_ == v) {
		left_ = u;
		u->parent_ = v->parent_;
	} else if (right_ == v) {
		right_ = u;
		u->parent_ = v->parent_;
	} else {
		std::cerr << "BSTNode::ReplaceChild Error: non-child passed as argument\n";
		exit(EXIT_FAILURE);
	}
}

AVLCommands::AVLCommands() : root_(nullptr), size_(0) {} // add height here?

int AVLCommands::max(int a, int b) {
	return (a > b)? a : b;
}

int AVLCommands::height(std::shared_ptr<BSTNode> currentNode) {
	if (root_ == nullptr) {
		return 0;
	}
	int leftHeight, rightHeight, heightMax;

	if (currentNode != nullptr) {
		if (currentNode->left_ != nullptr) {
			leftHeight = height(currentNode->left_);
		} else {
			leftHeight = -1;
		}

		if (currentNode->right_ != nullptr) {
			rightHeight = height(currentNode->right_);
		} else {
			rightHeight = -1;
		}

		heightMax = 1 + max(leftHeight, rightHeight);
		currentNode->height_ = heightMax;
		// update balance for each node
		currentNode->balance_ = leftHeight - rightHeight;
	}
  return currentNode->height_;
}

void AVLCommands::balance(std::shared_ptr<BSTNode> currentNode) {
	if (currentNode == nullptr) {
		return;
	}
	//std::cout << "balance called at key: " << currentNode->key_ << std::endl;
	std::shared_ptr<BSTNode> parent = currentNode->parent_.lock();
	if (currentNode != nullptr) {
		//std::cout << "Node's balance: " << currentNode->balance_ <<std::endl;

		if (currentNode->balance_ > 1) {
			// Left-Left rotation
			if (currentNode->left_->balance_ >= 0) {
				//std::cout << "Left-Left" << std::endl;
				rightRotate(currentNode);
				return;
				// Left-Right rotation
			} else if (currentNode->left_->balance_ < 0) {
				//std::cout << "Left-Right" << std::endl;
		    currentNode->left_ =  leftRotate(currentNode->left_);
		    rightRotate(currentNode);
				return;
			}
		} else if (currentNode->balance_ < -1) {
			// Right-Right rotation
			if (currentNode->right_->balance_ <= 0) {
				//std::cout << "Right-Right" << std::endl;
				leftRotate(currentNode);
				return;
				// Right-Left rotation
			} else if (currentNode->right_->balance_ > 0) {
				//std::cout << "Right-Left" << std::endl;
		    currentNode->right_ = rightRotate(currentNode->right_);
		    leftRotate(currentNode);
				return;
			}
		}
		balance(parent);
	}
	return;
}

std::shared_ptr<BSTNode> AVLCommands::rightRotate(std::shared_ptr<BSTNode> currentNode){
	if (currentNode == nullptr) {
		std::cout << "Can't rotate on an empty node!" << std::endl;
		return currentNode;
	}

	//std::cout << "Right Rotate at node " << currentNode->key_ << std::endl;
	std::shared_ptr<BSTNode> newSubRoot = currentNode->left_;
	currentNode->left_ = newSubRoot->right_;
	if (newSubRoot->right_ != nullptr) {
		newSubRoot->right_->parent_ = currentNode;
	}
	newSubRoot->right_ = currentNode;

	std::shared_ptr<BSTNode> oldParent = currentNode->parent_.lock();
	currentNode->parent_ = newSubRoot;

	if (oldParent == nullptr) {
		root_ = newSubRoot;
		newSubRoot->parent_.reset();
	} else {
		if (oldParent->right_== currentNode) {
			oldParent->right_ = newSubRoot;
		}else{
			oldParent->left_ = newSubRoot;
		}
		newSubRoot->parent_ = oldParent;
	}

	// Update heights and balances for the whole tree
	height(root_);
	// Return new root
	return newSubRoot;
}

// Takes the root of the tree to be left rotated and rotates it to provide new root
std::shared_ptr<BSTNode> AVLCommands::leftRotate(std::shared_ptr<BSTNode> currentNode){
	if (currentNode == nullptr) {
		std::cout << "Can't rotate on an empty node!" << std::endl;
		return currentNode;
	}

	//std::cout << "Left Rotate at node " << currentNode->key_ << std::endl;
	std::shared_ptr<BSTNode> newSubRoot = currentNode->right_;
	currentNode->right_ = newSubRoot->left_;
	if (newSubRoot->left_ != nullptr) {
		newSubRoot->left_->parent_ = currentNode;
	}
	newSubRoot->left_ = currentNode;

	std::shared_ptr<BSTNode> oldParent = currentNode->parent_.lock();
	currentNode->parent_ = newSubRoot;

	if (oldParent == nullptr) {
		root_ = newSubRoot;
		newSubRoot->parent_.reset();
	} else {
		if (oldParent->right_== currentNode) {
			oldParent->right_ = newSubRoot;
		}else{
			oldParent->left_ = newSubRoot;
		}
		newSubRoot->parent_ = oldParent;
	}

  //  Update heights and balances of the whole tree
	height(root_);
  // Return new root
  return newSubRoot;
}

// Basic tree printing for testing purposes
void AVLCommands::printTree(){
	printTree(root_);
}

void AVLCommands::printTree(std::shared_ptr<BSTNode> currentNode){
	if (root_ == nullptr) {
		std::cout << "tree is empty" << std::endl;
		return;
	} else {
		std::cout << "key: " << currentNode->key_ << std::endl;
		std::cout << "height: " << currentNode->height_ << std::endl;
		std::cout << "balance: " << currentNode->balance_ << std::endl;
		if (currentNode->parent_.lock() != nullptr) {
			std::cout << "parent: " << currentNode->parent_.lock()->key_ << std::endl;
		}
		if (currentNode->left_ != nullptr) {
			std::cout << "left" << std::endl;
			printTree(currentNode->left_);
		}
		if (currentNode->right_ != nullptr) {
			std::cout << "right" << std::endl;
			printTree(currentNode->right_);
		}
	}
}

void AVLCommands::Insert(int key) {
	//std::cout << "Inserting new key: " << key << std::endl;
	// BST insertion
	if (root_ == nullptr) {
		root_ = std::make_shared<BSTNode>(key);
		size_++;
		return;
	}
	std::shared_ptr<BSTNode> currentNode = root_, lastNode = nullptr;
	while (currentNode != nullptr) {
		lastNode = currentNode;
		currentNode = (key < currentNode->key_) ?
			currentNode->left_ : currentNode->right_;
	}
	if (key < lastNode->key_) {
		lastNode->left_ = std::make_shared<BSTNode>(key, lastNode);
		size_++;
	} else if (key > lastNode->key_) {
		lastNode->right_ = std::make_shared<BSTNode>(key, lastNode);
		size_++;
	}


	// update height and balance of the tree
	height(root_);
	balance(lastNode->parent_.lock());
}

bool AVLCommands::Delete(int key) {
	//std::cout << "Deleting key: " << key << std::endl;
	// base BST delete
	std::shared_ptr<BSTNode> currentNode = root_;
	std::shared_ptr<BSTNode> parent = nullptr;
	while (currentNode != nullptr) {
		if (currentNode->key_ == key) {
			if (currentNode->IsLeaf()) {
				DeleteLeaf(currentNode);

			} else if (currentNode->left_ == nullptr && currentNode->right_ != nullptr) {
				if (currentNode->parent_.lock() != nullptr) {
					std::shared_ptr<BSTNode> parent = currentNode->parent_.lock();
					parent->ReplaceChild(currentNode, currentNode->right_);
				} else { // if at the root
					root_ = currentNode->right_;
					currentNode = currentNode->right_;
					currentNode->parent_.reset();
					//std::cout << "Current node: " << currentNode->key_ << std::endl;
					//std::cout << "Right of current node: ";
					//std::cout << currentNode->right_->key_ << std::endl;
					//root_->ReplaceChild(currentNode->right_, currentNode);
				}
				size_--; assert(size_ >= 0);

			} else if (currentNode->right_ == nullptr && currentNode->left_ != nullptr) {
				if (currentNode->parent_.lock() != nullptr) {
					std::shared_ptr<BSTNode> parent = currentNode->parent_.lock();
					parent->ReplaceChild(currentNode, currentNode->left_);
				} else { // if at the root
					root_ = currentNode->left_;
					currentNode = currentNode->left_;
					currentNode->parent_.reset();
					//std::cout << "Current node: " << currentNode->key_ << std::endl;
					//std::cout << "Left of current node: ";
					//std::cout << currentNode->left_->key_ << std::endl;
					//root_->ReplaceChild(currentNode->left_, currentNode);
				}
				size_--; assert(size_ >= 0);

			} else {
				currentNode->key_ = DeleteMin(currentNode->right_);
			}
		}
		parent = currentNode->parent_.lock();
		currentNode = (key < currentNode->key_) ?
			currentNode->left_ : currentNode->right_;
	}

	// update height and balance of the tree
	height(root_);
	balance(parent);
	return false;
}


int AVLCommands::DeleteMin() {
	int deletedKey = DeleteMin(root_);
	// update the heights and balances of the whole tree
	height(root_);

	// find the smallest node after deletion
	std::shared_ptr<BSTNode> smallestNode = root_;
	std::shared_ptr<BSTNode> lastNode = nullptr;
	while (smallestNode != nullptr) {
		lastNode = smallestNode;
		smallestNode = smallestNode->left_;
	}
	// balance the tree
	balance(lastNode);
	return deletedKey;
}


void AVLCommands::DeleteLeaf(std::shared_ptr<BSTNode> currentNode) {
	std::shared_ptr<BSTNode> parent = currentNode->parent_.lock();
	if (parent == nullptr) {
		// Delete root
		root_ = nullptr;
		size_--; assert(size_ == 0);
	} else {
		if (parent->right_ == currentNode) {
			parent->right_ = nullptr;
		} else if (parent->left_ == currentNode) {
			parent->left_ = nullptr;
		} else {
			std::cerr << "BST::DeleteLeaf Error: inconsistent state\n";
		}
		size_--; assert(size_ >= 0);
	}
}

int AVLCommands::DeleteMin(std::shared_ptr<BSTNode> currentNode) {
	std::shared_ptr<BSTNode> lastNode = nullptr;
	while (currentNode != nullptr) {
		lastNode = currentNode;
		currentNode = currentNode->left_;
	}
	int result = lastNode->key_;
	std::shared_ptr<BSTNode> parent = lastNode->parent_.lock();
	if (parent == nullptr) {
		// lastNode is root
		if (lastNode->right_ != nullptr) {
			root_ = lastNode->right_;
			lastNode->right_->parent_.reset();
		} else {
			root_ = nullptr;
		}
	} else {
		// lastNode under the root
		if (lastNode->right_ != nullptr) {
			if (parent->left_ == lastNode){
				parent->left_ = lastNode->right_;
				lastNode->right_->parent_ = parent;
			} else {
				parent->right_ = lastNode->right_;
				lastNode->right_->parent_ = parent;
			}
		} else {
			if (parent->left_ == lastNode){
				parent->left_ = nullptr;
			} else {
				parent->right_ = nullptr;
			}
		}
  }
	size_--; assert(size_ >= 0);
	return result;
}

size_t AVLCommands::size() const {
	return size_;
}

bool AVLCommands::empty() const {
	return size_ == 0;
}

bool AVLCommands::Find(int key) const {
	std::shared_ptr<BSTNode> currentNode = root_;
	while (currentNode != nullptr) {
		if (currentNode->key_ == key) {
			return true;
		}
		currentNode = (key < currentNode->key_) ?
			currentNode->left_ : currentNode->right_;
	}
	return false;
}

std::string AVLCommands::JSON() const {
	nlohmann::json result;
	std::queue< std::shared_ptr<BSTNode> > nodes;
	if (root_ != nullptr) {
		result["height"] = root_->height_;
		result["root"] = root_->key_;
		nodes.push(root_);
		while (!nodes.empty()) {
			auto v = nodes.front();
			nodes.pop();
			std::string key = std::to_string(v->key_);
			std::string balance = std::to_string(v->balance_);
			std::string height = std::to_string(v->height_);
			result[key]["balance factor"] = balance;
			result[key]["height"] = height;
			if (v->left_ != nullptr) {
				result[key]["left"] = v->left_->key_;
				nodes.push(v->left_);
			}
			if (v->right_ != nullptr) {
				result[key]["right"] = v->right_->key_;
				nodes.push(v->right_);
			}
			if (v->parent_.lock() != nullptr) {
				result[key]["parent"] = v->parent_.lock()->key_;
			} else {
				result[key]["root"] = true;
			}
		}
	}
	result["size"] = size_;
	//std::cout << result << std::endl;
	return result.dump(2) + "\n";
}