Skip to content

Instantly share code, notes, and snippets.

@JonnoFTW
Created November 12, 2018 08:43
Show Gist options
  • Save JonnoFTW/76fd256c805f2b153f2bce219ccf00ef to your computer and use it in GitHub Desktop.
Save JonnoFTW/76fd256c805f2b153f2bce219ccf00ef to your computer and use it in GitHub Desktop.
from pydotplus import graphviz as pydot
class TrieNode:
__slots__ = ('text', 'children', 'is_end')
def __init__(self, letter, is_end=False):
self.text = letter
self.children = {}
self.is_end = is_end
def __getitem__(self, item):
if item in self.children:
return self.children[item]
else:
child = TrieNode(item)
self.children[item] = child
return child
def __repr__(self):
return "TrieNode: {} end={}".format(self.text if self.text else "root", self.is_end)
class Trie:
def __init__(self):
self.root = TrieNode("")
def add(self, word):
node = self.root
for letter in word:
node = node[letter]
node.is_end = True
def __contains__(self, word):
node = self.root
for letter in word:
node = node.children.get(letter, None)
if node is None:
return False
return node.is_end
def compress(self):
for l, child in self.root.children.items():
self._compress(self.root, child, l)
def _compress(self, parent, child, text):
if len(child.children) == 1 and len(parent.children) == 1:
parent.text += child.text
parent.is_end = child.is_end
# parent.children = child.children
else:
for l, sub_child in child.children.items():
self._compress(child, sub_child, l)
trie = Trie()
trie.add("I don't think so")
trie.add("I don't think the system works")
trie.add("I don't like sand")
trie.add("I see through the lies of the Jedi")
trie.add("I love democracy")
# trie.add("I sense Count Dooku")
# trie.add("It does not exist")
trie.compress()
tests = ["I hate you", "I love democracy"]
for line in tests:
print(line, "in trie:\t", line in trie)
dot = pydot.Dot(graph_type='graph', graph_name='Trie')
def plotTrie(node, plotParent, label=""):
for t, n in node.children.items():
pname = label + n.text
args = {'name': pname, 'color': 'red', 'label': t}
if n.is_end:
args['style'] = 'filled'
args['fillcolor'] = 'blue'
cdnode = pydot.Node(**args)
dot.add_edge(pydot.Edge(plotParent, cdnode))
plotTrie(n, cdnode, pname)
# once you've generated a compressed Trie called trie
i = 0
pdnode = pydot.Node('I')
plotTrie(trie.root["I"], pdnode)
dot.write_dot('memes.dot')
dot.write_png('memes.png')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment