Node.py

from abc import ABC, abstractmethod


class Node(ABC):
    @abstractmethod
    def predict(self, x):
        pass

    @abstractmethod
    def accept_visitor(self, visitor):
        pass


class Parent(Node):
    def __init__(self, feature_index, threshold, depth):
        self.feature_index = feature_index
        self.threshold = threshold
        self.depth = depth
        self.left_child = None
        self.right_child = None

    def predict(self, x):
        if (x[self.feature_index] < self.threshold):
            return self.left_child.predict(x)
        else:
            return self.right_child.predict(x)

    def accept_visitor(self, visitor):
        visitor.visit_parent(self)



class Leaf(Node):
    def __init__(self, value):
        self.value = value

    def predict(self, x):
        return self.value

    def accept_visitor(self, visitor):
        visitor.visit_leaf(self)