NominalInnerDecisionNode.h

#ifndef GRAIL_NOMINAL_INNER_DECISION_NODE_H
#define GRAIL_NOMINAL_INNER_DECISION_NODE_H
 
#include "DecisionNode.h"
#include <unordered_map>
#include <memory>
#include <iostream>
 
namespace grail
{
 namespace simulation
 { 
 template <class TDecisionType>
 class NominalInnerDecisionNode : public DecisionNode<TDecisionType>
 {
 public:
 
 NominalInnerDecisionNode(int column) :
 children{},
 columnIndex{ column }
 {
 
 }
 
 void SetDefaultDecision(std::unique_ptr< TDecisionType> decision)
 {
 defaultDecision = std::move(decision);
 }
 
 void AddChild(std::unique_ptr<DecisionNode<TDecisionType>> node, float key)
 {
 children[key] = std::move(node);
 }
 
 void Print(std::unordered_map<int, std::string>& columnNames, const std::string& indent) const override
 {
 if (columnNames.find(columnIndex) == columnNames.end())
 columnNames[columnIndex] = "Column[" + std::to_string(columnIndex) + "]";
 
 for(auto& child : children)
 {
 std::cout << indent << columnNames.at(columnIndex) << " = " << child.first << std::endl;
 child.second->Print(columnNames, indent + " ");
 }
 }
 
 const TDecisionType* Predict(std::vector<float>& data) const final
 {
 auto it = children.find(data[columnIndex]);
 if (it != children.end())
 {
 return it->second->Predict(data);
 }
 return defaultDecision.get();
 }
 
 DecisionNodeType GetNodeType() const final
 {
 return DecisionNodeType::NOMINAL_INNER;
 }
 private:
 std::unordered_map<float, std::unique_ptr<DecisionNode<TDecisionType>>> children;
 int columnIndex;
 std::unique_ptr<TDecisionType> defaultDecision = nullptr;
 
 template <class T> friend struct IDecisionTreeSerializer;
 };
 }
}
 
#endif //GRAIL_NOMINAL_INNER_DECISION_NODE_H