trie.py

class trie:
    def __init__(self):
        self.parentNodes = {}

    def search(self, word):
        if len(word) == 0 or not word[0] in self.parentNodes:
            return 0

        return self.parentNodes[word[0]].find(word[1:])

    def insert(self, word):
        if not word[0] in self.parentNodes:
            self.parentNodes[word[0]] = trieNode()

        self.parentNodes[word[0]].addChildren(word[1:])

class trieNode:
    def __init__(self):
        self.children = {}
        self.isAWord = False
        self.wordsBelow = 0

    def addChildren(self, word):
        if len(word) == 0 and not self.isAWord:
            self.isAWord = True
            return True
        elif len(word) == 0:
            return False

        if not word[0] in self.children:
            self.children[word[0]] = trieNode()

        if self.children[word[0]].addChildren(word[1:]):
            self.wordsBelow += 1
            return True
        else:
            return False

    def find(self, word):
        if len(word) == 0:
            return self.wordsBelow + 1 if self.isAWord else self.wordsBelow

        if not word[0] in self.children:
            return 1 if self.isAWord else 0

        return self.children[word[0]].find(word[1:])