Last active
August 15, 2018 16:49
-
-
Save rafiamafia/0f5a0fd1f1b177bd52e14579519e4234 to your computer and use it in GitHub Desktop.
A Trie class in python
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
class TrieNode: | |
def __init__(self): | |
self.children = dict() | |
self.end_of_word = False | |
def leafNode(self): | |
return self.end_of_word | |
def freeNode(self): | |
return not any(self.children) | |
class Trie: | |
def __init__(self): | |
self.root = TrieNode() | |
def insert(self, key): | |
current = self.root | |
for elem in key: | |
node = current.children.get(elem) | |
if node is None: | |
node = TrieNode() | |
current.children[elem] = node | |
current = node | |
current.end_of_word = True | |
def find(self, key): | |
current = self.root | |
for elem in key: | |
node = current.children.get(elem) | |
if node is None: | |
return False | |
current = node | |
return current.end_of_word | |
def delete(self, key): | |
if not self.find(key): | |
raise KeyError('{} does not exist.'.format(key)) | |
self._recursiveDelete(self.root, key, 0, len(key)) | |
# 1. Key present as a unique key i.e no part of key contains another key (prefix), not key itself is prefix to another key in trie. Delete all nodes. | |
# 2. Key present as a prefix key to another long key in trie. Unmark the leaf node. | |
# 3. Key present having atleast one other key as prefix key. Delete nodes from the end until first leaf node of longest prefix. | |
def _recursiveDelete(self, node, key, level, length): | |
if level == length: | |
if node.end_of_word: | |
node.end_of_word = False | |
return node.freeNode() # if node has no children it can be deleted | |
else: | |
child_node = node.children[key[level]] | |
if self._recursiveDelete(child_node, key, level+1, length): | |
del node.children[key[level]] | |
# recursively climb up and delete eligible nodes i.e can not be a leaf node or node has no children | |
return (not child_node.leafNode() and child_node.freeNode()) | |
return False |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment