This notebook was prepared by Donne Martin. Source and license info is on GitHub.

Solution Notebook¶

Constraints¶

• Can we assume we are working with strings?
• Yes
• Are the strings in ASCII?
• Yes
• Should find only match exact words with a terminating character?
• Yes
• Should list_words only return words with a terminating character?
• Yes
• Can we assume this fits memory?
• Yes

Test Cases¶

root node is denoted by ''

''
/  |  \
h   a*  m
/ \   \   \
a   e*  t*  e*
/ \         / \
s*  t*      n*  t*
/
s*

find

* Find on an empty trie
* Find non-matching
* Find matching

insert

* Insert on empty trie
* Insert to make a leaf terminator char
* Insert to extend an existing terminator char

remove

* Remove me
* Remove mens
* Remove a
* Remove has

list_words

* List empty
* List general case


Algorithm¶

find¶

• Set node to the root
• For each char in the input word
• Check the current node's children to see if it contains the char
• If a child has the char, set node to the child
• Else, return None
• Return the last child node if it has a terminator, else None

Complexity:

• Time: O(m), where m is the length of the word
• Space: O(h) for the recursion depth (tree height), or O(1) if using an iterative approach

insert¶

• set node to the root
• For each char in the input word
• Check the current node's children to see if it contains the char
• If a child has the char, set node to the child
• Else, insert a new node with the char
• Update children and parents
• Set the last node as a terminating node

Complexity:

• Time: O(m), where m is the length of the word
• Space: O(h) for the recursion depth (tree height), or O(1) if using an iterative approach

remove¶

• Determine the matching terminating node by calling the find method
• If the matching node has children, remove the terminator to prevent orphaning its children
• Set the parent node to the matching node's parent
• We'll be looping up the parent chain to propagate the delete
• While the parent is valid
• If the node has children
• If the node is a terminating node and it isn't the original matching node from the find call
• Remove the parent node's child entry matching the node
• Set the node to the parent
• Set the parent to the parent's parent

Complexity:

• Time: O(m+h), where where m is the length of the word and h is the tree height
• Space: O(h) for the recursion depth (tree height), or O(1) if using an iterative approach

list_words¶

• Do a pre-order traversal, passing down the current word
• When you reach a terminating node, add it to the list of results

Complexity:

• Time: O(n)
• Space: O(h) for the recursion depth (tree height), or O(1) if using an iterative approach

Code¶

In [1]:
%%writefile 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)

Overwriting trie.py

In [2]:
%run trie.py


Unit Test¶

In [3]:
%%writefile test_trie.py
import unittest

class TestTrie(unittest.TestCase):

def test_trie(self):
trie = Trie()

print('Test: Insert')
words = ['a', 'at', 'has', 'hat', 'he',
'me', 'men', 'mens', 'met']
for word in words:
trie.insert(word)
for word in trie.list_words():
self.assertTrue(trie.find(word) is not None)

print('Test: Remove me')
trie.remove('me')
words_removed = ['me']
words = ['a', 'at', 'has', 'hat', 'he',
'men', 'mens', 'met']
for word in words:
self.assertTrue(trie.find(word) is not None)
for word in words_removed:
self.assertTrue(trie.find(word) is None)

print('Test: Remove mens')
trie.remove('mens')
words_removed = ['me', 'mens']
words = ['a', 'at', 'has', 'hat', 'he',
'men', 'met']
for word in words:
self.assertTrue(trie.find(word) is not None)
for word in words_removed:
self.assertTrue(trie.find(word) is None)

print('Test: Remove a')
trie.remove('a')
words_removed = ['a', 'me', 'mens']
words = ['at', 'has', 'hat', 'he',
'men', 'met']
for word in words:
self.assertTrue(trie.find(word) is not None)
for word in words_removed:
self.assertTrue(trie.find(word) is None)

print('Test: Remove has')
trie.remove('has')
words_removed = ['a', 'has', 'me', 'mens']
words = ['at', 'hat', 'he',
'men', 'met']
for word in words:
self.assertTrue(trie.find(word) is not None)
for word in words_removed:
self.assertTrue(trie.find(word) is None)

print('Success: test_trie')

def test_trie_remove_invalid(self):
print('Test: Remove from empty trie')
trie = Trie()
self.assertTrue(trie.remove('foo') is None)

def main():
test = TestTrie()
test.test_trie()
test.assertRaises(KeyError, test.test_trie_remove_invalid)

if __name__ == '__main__':
main()

Overwriting test_trie.py

In [4]:
%run -i test_trie.py

Test: Insert
Test: Remove me
Test: Remove mens
Test: Remove a
Test: Remove has
Success: test_trie
Test: Remove from empty trie