interactive-coding-challenges/graphs_trees/trie/trie.py
2017-03-30 05:38:16 -04:00

76 lines
2.3 KiB
Python

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)