BasicTree

basic_tree.py

class Node:
    def __init__(self, value, count=0, children={}) -> None:
        self.value: str = value
        self.count: int = count
        self.children: dict[str, Node] = children

    def add_child(self, child) -> None:
        self.children[child.value] = child

    def __str__(self) -> str:
        return f"({self.value} [{self.count}])"
      

def print_tree(node, level=0):
  print("  " * level + f"{node.value} (count: {node.count})")
  for child in node.children.values():
      print_tree(child, level + 1)

if __name__ == "__main__":
    example_text = """
    Harry Potter is a wizard. He goes to Hogwarts. He is in Gryffindor. He has a friend named Ron.
    Ron has a rat named Scabbers. Harry has a pet owl named Hedwig.
    Harry has a wand. He has a scar on his forehead.
    """.strip()
    
    root = Node("root", 0, {})
    # Tokenize text
    tokenized_text = [sent.split() for sent in example_text.split(".")]
    # Create the tree
    for sent in tokenized_text:
        if not sent:
            continue
        current_node = root
        for word in sent:
            word = word.strip()
            if word in current_node.children:
                current_node = current_node.children[word]
                current_node.count += 1
            else:
                new_node = Node(word, 1, {})
                current_node.add_child(new_node)
                current_node = new_node

    print_tree(root)