Python script for visualising a Trie

safesparrow · safesparrow · commit b8d45e050a07 · 2022-10-30T23:09:05.000Z
diff --git a/input.txt b/input.txt
@@ -0,0 +1,3 @@
+abbrevia
+baba
+baba.ba
diff --git a/trie_vis.py b/trie_vis.py
@@ -0,0 +1,69 @@
+'''
+    A small program to visualize a trie data structure in the form of decision tree graph.
+'''
+#imports
+import json
+import pydot
+
+# function to return a new dict template
+def struct():
+    struct = {
+        
+    }
+    return struct
+
+# getting list of words as input from the file
+file_ = open('./input.txt', 'r')
+file_text = file_.read()
+file_len = len(file_text)
+file_.close()
+
+# trie making stuff happening (hard to explain)
+tmp_s = struct()
+root = tmp_s
+for line in file_text.split('\n'):
+    for c in line.split('.'):
+        if c not in tmp_s:
+            tmp_s[c] = struct()
+        tmp_s = tmp_s[c]
+    tmp_s = root
+    cur_word = []
+
+# saving the trie in a json file
+with open('output.json', 'w') as fp:
+    json.dump(root, fp, indent=4)
+
+# converting and saving the trie to dot language decision tree graph using pydot
+# // code taken and modified from stackoverflow (https://stackoverflow.com/questions/13688410/dictionary-object-to-decision-tree-in-pydot)
+rt = {'': root}
+counter = 0
+def draw(parent_name, child_name):
+    global counter
+    counter += 1
+    p_n = parent_name
+    c_n = child_name
+    graph.add_node(pydot.Node(p_n, label=parent_name.split('_')[0]))
+    graph.add_node(pydot.Node(c_n, label=child_name.split('_')[0]))
+    edge = pydot.Edge(p_n, c_n)
+    graph.add_edge(edge)
+
+def visit(node, parent=None):
+    global counter
+    for k,v in node.items():
+        if isinstance(v, dict):
+            # We start with the root node whose parent is None
+            # we don't want to graph the None node
+            k = k + '_' + str(counter)
+            if parent:
+                draw(parent, k)
+            visit(v, k)
+        else:
+            # drawing the label using a distinct name
+            v = v + '_' + str(counter)
+            draw(parent, v)
+
+graph = pydot.Dot(graph_type='digraph')
+visit(rt)
+graph.write_pdf('output.pdf')
+
+# end of program
