nicoloddo/x-tree-search

1from src.structures.markov_chain import MarkovChain2 3class Tree(MarkovChain):4    """5    Manages the tree structure including the root and the ability to expand the tree from each node.6    The nodes names are ids that refer to the path needed to reach them.7    The root node has id=0.8    A subsequent node would be for example:9    "0102" which means that to get there you need: root.children[1].children[0].children[2]10    """11    def __init__(self):12        super().__init__()13        self.root = self.__add_node() # Add root node14 15    class TreeNode(MarkovChain.MarkovNode):16        """17        Represents a single node in a tree.18 19        Attributes:20            parent (TreeNode): The TreeNode parent of this node. If it is None, this is the root node, identified with the id:"0".21            is_leaf (bool): Indicates whether this node is a leaf of the tree.22            value (object): Any object can be assigned to a node value.23            children (list) (derived): A list of Node instances that are the children of this node (derived property from the original MarkovNode.connections attribute)24            children_and_probs (list) (derived): List of tuples containing children of this node and the respective transition probability25        """26        def __init__(self, parent, value):27            self.parent = parent28            self.value = value    29            self.is_leaf = True30            self.children_and_probs = []31            self.children = []32 33            if parent == None:34                node_id = "0"35            else:36                node_id = parent.id + '_' + str(len(parent.children))37 38            super().__init__(node_id)39        40        def _add_child(self, child, probability):41            self._MarkovNode__add_connection(child, probability)42            if len(self.connections) > 0:43                self.is_leaf = False44            self.update_children_and_probs()45        46        def update_children_and_probs(self):47            """Updates the list of tuples containing children of this node and the respective transition probability"""48            self.children_and_probs = self.connections.items()49            self.update_children()50 51        def update_children(self):52            """Updates the list of children nodes of this node"""53            self.children = [t[0] for t in self.children_and_probs] # Extract the first item of the tuple, which is the pointer to the child node54 55        @property56        def id(self):57            return self.name58 59        @id.setter60        def id(self, value):61            self.name = value62 63    def __add_node(self, parent=None, value=None):64        """Overrides the add_node method to ensure TreeNode objects are created."""65        new_node = self.TreeNode(parent, value)66        self.nodes[new_node.id] = new_node67        return new_node68 69    def set_children(self, node_id, children_values):70        """71        Expands a node with a given amount of children and their values72 73        Args:74            node (str): The id of the node to expand with new children.75            children_values (list of numbers): Values to assign to the children nodes.76        """77        parent = self.nodes[node_id]78        for value in children_values:79            child = self.__add_node(parent=parent, value=value)80            probability = 1/len(children_values)81            parent._add_child(child, probability)82 83    def get_node(self, node_id):84        """85        Returns a node given its id.86 87        Args:88            node (str): The id of the node to get.89        """90        return self.nodes[node_id]91 92    def __str__(self):93        result = "Tree\n"94        for node_id, node in self.nodes.items():95            result += f"\nNode: {node_id}\nChildren:\n"96            for child, prob in node.children_and_probs:97                result += f"{child.id} with P = {prob.item()} and value V = {child.value}\n" # remember that prob is a tensor98        return result99 100 101 102        103 104