interactive-coding-challenges/graphs_trees/trie/trie.py

from collections import OrderedDict


class Node(object):

    def __init__(self, key, parent=None, terminates=False):
        self.key = key
        self.terminates = False
        self.parent = parent
        self.children = {}


class Trie(object):

    def __init__(self):
        self.root = Node('')

    def find(self, word):
        if word is None:
            raise TypeError('word cannot be None')
        node = self.root
        for char in word:
            if char in node.children:
                node = node.children[char]
            else:
                return None
        return node if node.terminates else None

    def insert(self, word):
        if word is None:
            raise TypeError('word cannot be None')
        node = self.root
        parent = None
        for char in word:
            if char in node.children:
                node = node.children[char]
            else:
                node.children[char] = Node(char, parent=node)
                node = node.children[char]
        node.terminates = True

    def remove(self, word):
        if word is None:
            raise TypeError('word cannot be None')
        node = self.find(word)
        if node is None:
            raise KeyError('word does not exist')
        node.terminates = False
        parent = node.parent
        while parent is not None:
            # As we are propagating the delete up the 
            # parents, if this node has children, stop
            # here to prevent orphaning its children.
            # Or
            # if this node is a terminating node that is
            # not the terminating node of the input word, 
            # stop to prevent removing the associated word.
            if node.children or node.terminates:
                return
            del parent.children[node.key]
            node = parent
            parent = parent.parent

    def list_words(self):
        result = []
        curr_word = ''
        self._list_words(self.root, curr_word, result)
        return result

    def _list_words(self, node, curr_word, result):
        if node is None:
            return
        for key, child in node.children.items():
            if child.terminates:
                result.append(curr_word+key)
            self._list_words(child, curr_word+key, result)