Home > Enterprise >  Store tree as a dictionary using its root
Store tree as a dictionary using its root

Time:11-26

I have a root node with attributes data, parent, and children (list). I want to use a method to recursively store the entire tree in a dictionary.

treeNode Class

class treeNode:
    is_leaf = True
    root = False
    
    def __init__(self,data):
        self.data = data
        self.parent = None
        self.children = []

    def __repr__(self):
        return f"{self.data}"
    
    def add_child(self,child):
        child.parent = self
        self.children.append(child)
        self.is_leaf = False
    
    def get_parent(self):
        if self.parent is not None:
            return self.parent
        else:
            self.root = True
            return None
    
    def get_children(self):
        if not self.is_leaf:
            return self.children
        else:
            print("This is a leaf")

    def set_frequency(self,freq):
        self.frequency = freq

    ##This is the method I want to use
    ## First input (node) is root
    def get_all_tree(self,node):       
        if node.is_leaf:
            children.append(node)
        else:
            for x in node.get_children():
                self.get_all_tree(x) 

How do I implement this?

Edit: I want to create a dictionary tree that is somewhat like this:

{0: [1, {2: [11, 12, 13, 16, {3: [14, 15]}, {4: [41, 42, 43]}]}]}

CodePudding user response:

class treeNode:
...
    def get_all_tree(self, node):
        if node.is_leaf:
            return node
        else:
            tree = {node: []}
            for x in node.get_children():
                tree[node].append(x.get_all_tree(x))
            return tree

root = treeNode('0')
child1 = treeNode('1')
child2 = treeNode('2')
root.add_child(child1)
root.add_child(child2)
child1_1 = treeNode('1.1')
child1_2 = treeNode('1.2')
child1.add_child(child1_1)
child1.add_child(child1_2)
child1_1_1 = treeNode('1.1.1')
child1_1.add_child(child1_1_1)
print(root.get_all_tree(root))

This returns {0: [{1: [{1.1: [1.1.1]}, 1.2]}, 2]}

There's a few improvements to be made in your code I believe, as you can see, it is quite tedious to add children. Considering the class is named treeNode I believe you also have a Tree class, I recommend moving the handling of children and parents in the Tree class but keep the information about each node parent/children within the treeNode class. I will try to come back later with a minimal example and explaining it.

Good job on your implementation so far though !

  • Related