[Data Structure] Tree C++로 만들기

#include <iostream>
#include <string>
#include <list>
#include <vector>
#include <cmath>
using namespace std;

struct Node
{
	int value;
	int depth;
	Node* child[2];                     // 0 - left, 1 - right
	int sibiling;						// 같은 depth에서 좌측부터 부여된 index

	Node()
	{
		depth = 0;
		value = 0;
		sibiling = 0;
		child[0] = child[1] = NULL;
	}
};

Node* root = NULL;

bool valueComp(Node* i, Node* j)
{
	return	i->value < j->value;
}

bool DepthComp(Node* i, Node* j)
{
	return	i->depth < j->depth;
}

void CollectTreeNode(std::list<Node*>& listNode, Node* node, Node* parent, int childIndex)
{
	if (parent)
	{
		node->sibiling = parent->sibiling * 2 + childIndex;
	}
	else
	{
		node->sibiling = 0;
	}

	listNode.push_back(node);

	for (int i = 0; i < 2; ++i)
	{
		if (node->child[i])
		{
			CollectTreeNode(listNode, node->child[i], node, i);
		}
	}
}

void ConstructTree(std::list<Node*>& listNode, Node* parent, int childIndex)
{
	int listLen = listNode.size();

	int halfIndex = listLen / 2;
	Node* nwCenter = NULL;
	std::list<Node*> listLeft;
	std::list<Node*> listRight;

	int i = 0;
	std::list<Node*>::iterator itor_node = listNode.begin();
	while (itor_node != listNode.end())
	{
		Node* n = *itor_node;
		if (i < halfIndex)
		{
			listLeft.push_back(*itor_node);
		}
		else if (i == halfIndex)
		{
			nwCenter = *itor_node;
			if (parent == NULL)
			{
				root = nwCenter;
				nwCenter->depth = 0;
			}
			else
			{
				int index = nwCenter->value > parent->value ? 1 : 0;
				int index2 = 1 - index;
				parent->child[index] = nwCenter;
				nwCenter->depth = parent->depth + 1;
			}
		}
		else
		{
			listRight.push_back(*itor_node);
		}

		++i;
		++itor_node;
	}

	if (listLeft.size() > 0)
	{
		ConstructTree(listLeft, nwCenter, 0);
	}
	if (listRight.size() > 0)
	{
		ConstructTree(listRight, nwCenter, 1);
	}
}

void RearrangeTree()
{
	std::list<Node*> listNode;
	CollectTreeNode(listNode, root, NULL, 0);
	listNode.sort(valueComp);

	list<Node*>::iterator itr = listNode.begin();
	while (itr != listNode.end())
	{
		Node* node = *itr;
		node->child[0] = NULL;
		node->child[1] = NULL;
		itr++;
	}

	ConstructTree(listNode, NULL, 0);
}

void AddNodeToTree(Node* parent, Node* node)
{
	if (parent == NULL)
	{
		root = node;
		return;
	}

	int addIndex = node->value > parent->value ? 1 : 0;
	if (parent->child[addIndex] == NULL)
	{
		parent->child[addIndex] = node;
		node->depth = parent->depth + 1;
		RearrangeTree();
	}
	else
	{
		AddNodeToTree(parent->child[addIndex], node);
	}
}

void RemoveNodeFromTree(int val, Node* parent, int childIndex, Node* node)
{
	if (node == NULL) return;

	if (node->value == val)
	{
		if (node->child[0] == NULL && node->child[1] == NULL)
		{
			if (parent == NULL)
			{
				root = NULL;
				return;
			}
			else
			{
				parent->child[childIndex] = NULL;
				// tree depth 갱신

				return;
			}
		}
		else
		{
			std::list<Node*> listNode;
			CollectTreeNode(listNode, root, NULL, 0);
			
			list<Node*>::iterator iter = listNode.begin();
			while (iter != listNode.end())
			{
				Node* pNode = *iter;
				if (pNode->value == val)
				{
					listNode.erase(iter);
					break;
				}
				++iter;
			}

			iter = listNode.begin();
			while (iter != listNode.end())
			{
				Node* node = *iter;
				node->child[0] = NULL;
				node->child[1] = NULL;
				iter++;
			}

			listNode.sort(valueComp);
			ConstructTree(listNode, NULL, 0);
		}
	}
	else
	{
		int index = node->value < val ? 1 : 0;
		RemoveNodeFromTree(val, node, index, node->child[index]);
	}
}

Node* FindNodeByValue(int val, Node* node)
{
	if (node->value == val)
	{
		return node;
	}

	int index = node->value < val ? 1 : 0;
	Node* childNode = node->child[index];
	if (childNode == NULL)
	{
		return NULL;
	}
	else
	{
		return FindNodeByValue(val, childNode);
	}
}

void RemoveNodeByKeyFromTree(int val)
{
	RemoveNodeFromTree(val, NULL, 0, root);
}

Node* FindParentNode(Node* childNode, Node* parentNode)
{
	int index = childNode->value > parentNode->value ? 1 : 0;
	if (parentNode->child[index] == childNode)
	{
		return parentNode;
	}
	if (childNode == parentNode)
	{
		return NULL;
	}
	FindParentNode(childNode, parentNode->child[index]);
}

void DisplayTree()
{
	std::list<Node*> listNode;
	CollectTreeNode(listNode, root, NULL, 0);
	listNode.sort(DepthComp);

	int maxDepth = listNode.back()->depth;
	vector<vector<Node*>> vecDepthVectors(maxDepth+1);
	//vecDepthVectors.reserve(maxDepth);
	for (int i = 0; i < maxDepth+1; ++i)
	{
		size_t capacity = pow(2, i);
		vecDepthVectors[i].reserve(capacity);
		vecDepthVectors[i].assign(capacity, NULL);

	}

	std::list<Node*>::iterator iter = listNode.begin();
	while (iter != listNode.end())
	{
		Node* pNode = *iter;
		int sibling = pNode->sibiling;
		vector<Node*>::iterator itr = (vecDepthVectors[pNode->depth].begin() + sibling);
		vecDepthVectors[pNode->depth].insert(itr, pNode);
		++iter;
	}

	printf("\n[Root]\n");
	

	for (int depth = 0; depth <= maxDepth; ++depth)
	{
		int maxSibling = pow(2, depth);
		// 시작 space 출력
		int tab = maxDepth - depth;
		for (int i = 0; i < tab; i++)
		{
			printf("   ");
		}

		for (int sibiling = 0; sibiling < maxSibling; ++sibiling)
		{
			Node* pNode = vecDepthVectors[depth][sibiling];
			if (pNode)
			{
				printf("[%03d]", pNode->value);		// 001
			}
			else
			{
				printf("   ");
			}

			if (sibiling % 2 == 0)
			{
				// sibilng 간 space 출력 (작게)
				for (int i = 0; i < tab; i++)
				{
					printf("   ");
				}
			}
			else
			{
				// sibilng 간 space 출력 (크게)
				for (int i = 0; i < tab; i++)
				{
					printf(" ");
				}

			}
		}
		printf("\n");
	}
}

int main()
{
	const int size = 20;
	printf("Size : %d\n", size);

	Node* nodes[size];
	for (int i = 0; i < size; ++i)
	{
		Node* N = new Node;
		N->value = i;
		AddNodeToTree(root, N);
		nodes[i] = N;
	}

	RemoveNodeByKeyFromTree(4);
	DisplayTree();

	RemoveNodeByKeyFromTree(10);
	DisplayTree();

	getchar();
}