TreeBranch.hpp

Go to the documentation of this file.

#pragma once
 
#include "models/TreeNode.hpp"
#include "models/Projector.hpp"
#include "utils/Types.hpp"
 
namespace ppforest2 {
  class TreeBranch final : public TreeNode {
  public:
    using Ptr = std::unique_ptr<TreeBranch>;

    pp::Projector projector;
    types::Feature cutpoint;
    TreeNode::Ptr lower;
    TreeNode::Ptr upper;

    std::set<types::GroupId> groups;
    types::Feature pp_index_value = 0;
 
    TreeBranch(
        pp::Projector projector,
        types::Feature cutpoint,
        TreeNode::Ptr lower,
        TreeNode::Ptr upper,
        std::set<types::GroupId> groups = {},
        types::Feature pp_index_value   = 0
    );
 
    void accept(TreeNode::Visitor& visitor) const override;

    types::Outcome response() const override;

    types::Outcome predict(types::FeatureVector const& x) const override;
 
    int group_count() const override { return static_cast<int>(groups.size()); }
 
    std::set<types::GroupId> node_groups() const override { return groups; }
 
    bool equals(TreeNode const& other) const override;
    TreeNode::Ptr clone() const override;

    static Ptr make(
        pp::Projector projector,
        types::Feature cutpoint,
        TreeNode::Ptr lower,
        TreeNode::Ptr upper,
        std::set<types::GroupId> groups = {},
        types::Feature pp_index_value   = 0
    );
  };
}