prefix_trie.py

from collections import deque

class TrieNode:
    def __init__(self, char):
        self.char = char
        self.is_word = False
        self.children = {}

class Trie:
    def __init__(self):
        self._root = TrieNode('')

    def insert(self, word, idx=0, node=None):
        if node is None:
            node = self._root

        ch = word[idx]
        if ch in node.children:
            self.insert(word, idx+1, node.children.get(ch))
        else:
            new_child = TrieNode(ch)
            node.children[ch] = new_child
            if idx == len(word)-1:
                new_child.is_word = True
            else:
                self.insert(word, idx+1, new_child)

    def starts_with(self, prefix, i=0, node=None):
        if not node: node = self._root
        ch = prefix[i]
        if ch not in node.children:
            return False
        if i == len(prefix)-1:
            return True
        else:
            return self.starts_with(prefix, i+1, node.children[ch])

    def find(self, word):
        found, _parent_node = self._search(word)
        return found

    def delete(self, word):
        found, parent_node = self._search(word)
        last_char = word[-1]
        if found and not parent_node.children[last_char].children:
            del parent_node.children[last_char]

    def _search(self, word, i=0, node=None):
        if not node: node = self._root
        ch = word[i]
        if ch not in node.children:
            return False, None
        if i == len(word)-1:
            return node.children[ch].is_word, node
        else:
            return self._search(word, i+1, node.children[ch])