We define Trie as a tree structure while each node uses a dict to store edges (We can also use a char array with 256 length for ASCII encoding, but usually it is not a complete tree so using a dict will potentially save some memory).
class TrieNode:
def __init__(self):
self.children = {}
self.isWord = False
Here we uses a flag: isWord
to indicates that is there a word ending with current node.
A trie a tree structure that contains many TrieNode
which supports insert
, find
, and findPrefix
. A special implementation may also support removal.
class Trie:
def __init__(self):
self.root = TrieNode()
def insert(self, word):
if word == '':
self.root.isWord = True
ptr = self.root
for c in word:
if c not in ptr.children:
ptr.children[c] = TrieNode()
ptr = ptr.children[c]
ptr.isWord = True
def find(self, word):
if word == '':
return self.root.isWord
ptr = self.root
for c in word:
if c in ptr.children:
ptr = ptr.children[c]
else:
return False
return ptr.isWord
def findPerfix(self, prefix):
if prefix == '':
return self.root.isWord
ptr = self.root
for c in prefix:
if c in ptr.children:
ptr = ptr.children[c]
else:
return False
return True
To support removal, we need to add another flag hasWord
within the TrieNode
which indicates if the node is still contains prefix. Because we do not physically remove a existing child from its parent.
class TrieNode:
def __init__(self):
self.children = {}
self.isWord = False
self.hasWord = False
class Trie:
def __init__(self):
self.root = TrieNode()
def insert(self, word):
"""
At each level we also updates hasWord to True
"""
if word == '':
self.root.isWord = True
self.root.hasWord = True
ptr = self.root
for c in word:
if c not in ptr.children:
ptr.children[c] = TrieNode()
ptr.hasWord = True
ptr = ptr.children[c]
ptr.isWord = True
ptr.hasWord = True
def find(self, word):
if word == '':
return self.root.isWord
ptr = self.root
for c in word:
# because the node may contains child that has been remove but still exist in children list, we need to first check whether it is still valid
if not ptr.hasWord:
return False
if c in ptr.children:
ptr = ptr.children[c]
else:
return False
return ptr.isWord
def findPerfix(self, prefix):
if prefix == '':
return self.root.isWord
ptr = self.root
for c in prefix:
if not ptr.hasWord:
return False
if c in ptr.children:
ptr = ptr.children[c]
else:
return False
# !!! note: check if it is still valid
return ptr.hasWord
def remove(self, word):
if word == '':
self.root.isWord = False
self.root.hasWord = False
def helper(node, key):
if key == '':
node.isWord = False
node.hasWord = False
return
# the word is not in the trie
if key[0] not in node.children:
return
helper(node.children[key[0]], key[1:])
# backtrack and update parent flag with children flags
node.hasWord = any([child.hasWord for child in node.children.values()])
helper(self.root, word)