trie.py

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')