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)