decisiontree.cpp

A simple decision tree c++ implementation

// DecisionTree.cpp: implementation of the DecisionTree class.
//
//////////////////////////////////////////////////////////////////////

#include "DecisionTree.h"
#include "TreeNode.h"
#include <iostream>


//////////////////////////////////////////////////////////////////////
// Construction/Destruction
//////////////////////////////////////////////////////////////////////

DecisionTree::DecisionTree()
{
	//set roots node to NULL upon tree creation
	m_pRootNode = NULL;
}

DecisionTree::~DecisionTree()
{
	//recursive call to delete tree nodes and root
	RemoveNode(m_pRootNode);
}

void DecisionTree::CreateRootNode(int nodeID, string newQorA)
{
	//create the root node with a specific ID and string
	m_pRootNode = new TreeNode(nodeID, newQorA);
}

void DecisionTree::AddYesNode(int existingNodeID, int newNodeID, string newQorA)
{
	//if you dont have a root node you cant add another node
	if(m_pRootNode == NULL)
	{
		cout << "Error - no root node in AddYesNode()\n";
		return;
	}

	//otherwise query tree and add node
	if(SearchTreeAndAddYesNode(m_pRootNode, existingNodeID, newNodeID, newQorA))
	{
		cout << "Added 'yes' node";
		cout << newNodeID;
		cout << " onto 'yes' branch of node ";
		cout << existingNodeID;
		cout << endl;
	}
	else
	{
		cout << "Node ";
		cout << existingNodeID;
		cout << " not found \n";
	}
}

bool DecisionTree::SearchTreeAndAddYesNode(TreeNode* currentNode, int existingNodeID, int newNodeID, string newQorA)
{
	
	if(currentNode->m_iNodeID == existingNodeID)
	{
		//create node
		if(currentNode->m_pYesBranch == NULL)
		{
			currentNode->m_pYesBranch = new TreeNode(newNodeID, newQorA);
		}
		else
		{
			currentNode->m_pYesBranch = new TreeNode(newNodeID, newQorA);
		}
		return true;
	}
	else
	{
		//try yes branch if it exists
		if(currentNode->m_pYesBranch != NULL)
		{
			if(SearchTreeAndAddYesNode(currentNode->m_pYesBranch, existingNodeID, newNodeID, newQorA))
			{
				return true;
			}
			else
			{
				//try no branch if it exists
				if(currentNode->m_pNoBranch != NULL)
				{
					return(SearchTreeAndAddYesNode(currentNode->m_pNoBranch, existingNodeID, newNodeID, newQorA));
				}
				else
					return false;
			}
		}
		return false;
	}
}

void DecisionTree::AddNoNode(int existingNodeID, int newNodeID, string newQorA)
{
	if(m_pRootNode == NULL)
	{
		cout << "Error no root node in AddNoNode()\n";
		return;
	}
	if(SearchTreeAndAddNoNode(m_pRootNode, existingNodeID, newNodeID, newQorA))
	{
		cout << "Added 'no' node";
		cout << newNodeID;
		cout << " onto 'no' branch of node ";
		cout << existingNodeID;
		cout << endl;
	}
	else
	{
		cout << "Node ";
		cout << existingNodeID;
		cout << " not found \n";
	}

}

bool DecisionTree::SearchTreeAndAddNoNode(TreeNode* currentNode, int existingNodeID, int newNodeID, string newQorA)
{
	if(currentNode->m_iNodeID == existingNodeID)
	{
		if(currentNode->m_pNoBranch == NULL)
		{
			currentNode->m_pNoBranch = new TreeNode(newNodeID, newQorA);
		}
		else
		{
			currentNode->m_pNoBranch = new TreeNode(newNodeID, newQorA);
		}
		return true;
	}
	else
	{
		if(currentNode->m_pYesBranch != NULL)
		{
			if(SearchTreeAndAddNoNode(currentNode->m_pYesBranch, existingNodeID, newNodeID, newQorA))
			{
				return true;
			}
			else
			{
				if(currentNode->m_pNoBranch != NULL)
				{
					return(SearchTreeAndAddNoNode(currentNode->m_pNoBranch, existingNodeID, newNodeID, newQorA));
				}
				else
					return false;

			}
		}
		else 
			return false;
	}
}

void DecisionTree::QueryBinaryTree(TreeNode* currentNode)
{
	if(currentNode->m_pYesBranch == NULL)
	{
		//if both the yes and no branch pointers are NULL 
		//the tree is at a decision outcome state so output
		//the string
		if(currentNode->m_pNoBranch == NULL)
			cout << currentNode->m_strQuestOrAns << endl;
		else
			cout << "Missing yes branch at " + currentNode->m_strQuestOrAns + " question\n";	
		return;
	}
	if(currentNode->m_pNoBranch == NULL)
	{
		cout << "Missing no branch at " + currentNode->m_strQuestOrAns + " question\n";
		return;
	}

	//otherwise default to asking the question at the currentNode
	AskQuestion(currentNode);
}

void DecisionTree::Query()
{
	QueryBinaryTree(m_pRootNode);
}

void DecisionTree::AskQuestion(TreeNode *node)
{
	cout << node->m_strQuestOrAns + " (enter yes or no)\n";
	string answer;
	cin >> answer;
	if(answer == "yes")
		QueryBinaryTree(node->m_pYesBranch);
	else if(answer == "no")
		QueryBinaryTree(node->m_pNoBranch);
	else
	{
		cout << "Error please answer yes or no\n";
		AskQuestion(node);
	}
}

void DecisionTree::Output()
{
	OutputBinaryTree("1", m_pRootNode);
}

void DecisionTree::OutputBinaryTree(string tag, TreeNode* currentNode)
{
	if(currentNode == NULL)
		return;

	cout << "[" + tag + "] node id = ";
	
	cout << currentNode->m_iNodeID;
	cout << ", question/answer = ";
	cout << currentNode->m_strQuestOrAns;
	cout << endl;
	
	// Go down yes branch
	OutputBinaryTree(tag + ".1",currentNode->m_pYesBranch);
	// Go down no branch
	OutputBinaryTree(tag + ".2",currentNode->m_pNoBranch);

}

void DecisionTree::RemoveNode(TreeNode *node)
{
	if(node != NULL)
	{
		if(node->m_pYesBranch != NULL)
			RemoveNode(node->m_pYesBranch);

		if(node->m_pNoBranch != NULL)
			RemoveNode(node->m_pNoBranch);
		
		cout << "deleting node " << node->m_iNodeID << endl;
		delete node;
		node = NULL;
	}
}