How to print Trie in tree structure in python?

Question

I am still a beginner for programming. I do not know how to print trie in tree structure as I have try for a few methods. Anyone can help?

class TrieNode:
 
    def __init__(self, data):
 
        self.data = data
        self.is_end = False
        self.children = {}
 
class Trie(object):
 
    def __init__(self):
        self.root = TrieNode("")

    def insert(self, array):
        node = self.root
 
        for x in array:
            if x in node.children:
                node = node.children[x]
                print(node.data)
            else:
                new_node = TrieNode(x)
                node.children[x] = new_node
                node = new_node
                print(node.data)
        node.is_end = True

Below is the code from the main class:-

tr = Trie()
n = int(input("Enter number of file(s): "))
for x in range (n):
    path = input("Enter your directory path: ")
    tr.insert(path.split("/"))
print(tr.root.data)

Answer 1

You could define a __repr__ method on your Trie class so that when you do print(tr) , that method will be called, and you'll get the tree printed in the way you want.

Here is possible implementation for such a __repr__ method: it uses a recursive function and an indent variable that will ensure that items are indented:

    def __repr__(self):
        def recur(node, indent):
            return "".join(indent + key + ("$" if child.is_end else "") 
                                  + recur(child, indent + "  ") 
                for key, child in node.children.items())

        return recur(self.root, "\n")

This string representation will append a $ after each data when the corresponding is_end attribute is True.

A remark about your trie implementation:

You are currently storing the data twice: once as a data attribute, and a second time as a key in the children dictionary. This should not be necessary. At the two places where you use node.data , you could use x instead -- they have the same value.

You can then remove the data attribute from your TrieNode class. The code will then look like this:

class TrieNode:
    def __init__(self):
        self.is_end = False
        self.children = {}


class Trie(object):
    def __init__(self):
        self.root = TrieNode()

    def insert(self, array):
        node = self.root
 
        for x in array:
            if x in node.children:
                node = node.children[x]
            else:
                child = TrieNode()
                node.children[x] = child
                node = child
            print(x)
        node.is_end = True

    def __repr__(self):
        def recur(node, indent):
            return "".join(indent + key + ("$" if child.is_end else "") 
                                  + recur(child, indent + "  ") 
                for key, child in node.children.items())

        return recur(self.root, "\n")


# Main code
tr = Trie()
n = int(input("Enter number of file(s): "))
for x in range (n):
    path = input("Enter your directory path: ")
    tr.insert(path.split("/"))
print(tr)