from typing import List, Dict
import json

node_count = 0

class Node:
    def __init__(self, name, childs:List["Node"]):
        global node_count
        self.name = (name, node_count)
        node_count += 1
        self.childs = childs

    def __repr__(self):
        return str({self.name: self.childs})


def make_tree(depth) -> Node:
    if depth <= 0:
        return Node(name=depth, childs=[])
    else:
        return Node(name=depth, childs=[make_tree(depth-1) for i in range(2)])


class TreeTraversalHelper:
    def __init__(self, tree:Node):
        self.current_stage = [tree]
        self.next_stage = []
        self.all = []

    def traverse(self):
        for node in self.current_stage:
            print(node.name)
            self.all.append(node)
            for child in node.childs:
                self.next_stage.append(child)

        if self.next_stage:
            self.current_stage = self.next_stage
            self.next_stage = []
            self.traverse()


if __name__ == "__main__":
    t = make_tree(3)
    T = TreeTraversalHelper(t)
    T.traverse()
    print([node.name for node in T.all])

By admin

Leave a Reply

Your email address will not be published.